diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py
index f8fda13a35..1f1db914e2 100755
--- a/onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py
+++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py
@@ -864,7 +864,7 @@ def run_tensorrt_static_xl(
     base_pipeline_info = PipelineInfo(version)
     demo_base = init_pipeline(Txt2ImgXLPipeline, base_pipeline_info)
 
-    refiner_pipeline_info = PipelineInfo(version, is_sd_xl_refiner=True)
+    refiner_pipeline_info = PipelineInfo(version, is_refiner=True)
     demo_refiner = init_pipeline(Img2ImgXLPipeline, refiner_pipeline_info)
 
     max_device_memory = max(demo_base.backend.max_device_memory(), demo_refiner.backend.max_device_memory())
@@ -887,7 +887,7 @@ def run_tensorrt_static_xl(
             guidance=5.0,
             warmup=warmup,
             seed=seed,
-            return_type="latents",
+            return_type="latent",
         )
 
         images, time_refiner = demo_refiner.run(
@@ -1037,7 +1037,7 @@ def run_ort_trt_xl(
     base_pipeline_info = PipelineInfo(version)
     demo_base = init_pipeline(Txt2ImgXLPipeline, base_pipeline_info)
 
-    refiner_pipeline_info = PipelineInfo(version, is_sd_xl_refiner=True)
+    refiner_pipeline_info = PipelineInfo(version, is_refiner=True)
     demo_refiner = init_pipeline(Img2ImgXLPipeline, refiner_pipeline_info)
 
     demo_base.load_resources(image_height, image_width, batch_size)
@@ -1053,7 +1053,7 @@ def run_ort_trt_xl(
             guidance=5.0,
             warmup=warmup,
             seed=seed,
-            return_type="latents",
+            return_type="latent",
         )
         images, time_refiner = demo_refiner.run(
             prompt,
diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_txt2img.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_txt2img.py
index f6e00063a6..d6de5c45a5 100644
--- a/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_txt2img.py
+++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_txt2img.py
@@ -73,7 +73,7 @@ if __name__ == "__main__":
             denoising_steps=args.denoising_steps,
             guidance=args.guidance,
             seed=args.seed,
-            return_type="images",
+            return_type="image",
         )
 
     if not args.disable_cuda_graph:
diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_txt2img_xl.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_txt2img_xl.py
index c3a2e4e293..efc87a207d 100644
--- a/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_txt2img_xl.py
+++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_txt2img_xl.py
@@ -28,11 +28,15 @@ from engine_builder import EngineType, get_engine_type
 from pipeline_img2img_xl import Img2ImgXLPipeline
 from pipeline_txt2img_xl import Txt2ImgXLPipeline
 
-if __name__ == "__main__":
-    coloredlogs.install(fmt="%(funcName)20s: %(message)s")
+
+def run_demo():
+    """Run Stable Diffusion XL Base + Refiner together (known as ensemble of expert denoisers) to generate an image."""
+
     args = parse_arguments(is_xl=True, description="Options for Stable Diffusion XL Demo")
+
     prompt, negative_prompt = repeat_prompt(args)
 
+    # Recommend image size as one of those used in training (see Appendix I in https://arxiv.org/pdf/2307.01952.pdf).
     image_height = args.height
     image_width = args.width
 
@@ -44,39 +48,32 @@ if __name__ == "__main__":
         init_trt_plugins()
 
     max_batch_size = 16
-    if args.build_dynamic_shape or image_height > 512 or image_width > 512:
+    if (engine_type in [EngineType.ORT_TRT, EngineType.TRT]) and (
+        args.build_dynamic_shape or image_height > 512 or image_width > 512
+    ):
         max_batch_size = 4
 
     batch_size = len(prompt)
     if batch_size > max_batch_size:
-        raise ValueError(
-            f"Batch size {len(prompt)} is larger than allowed {max_batch_size}. If dynamic shape is used, then maximum batch size is 4"
-        )
+        raise ValueError(f"Batch size {batch_size} is larger than allowed {max_batch_size}.")
 
-    base_info = PipelineInfo(args.version, use_vae_in_xl_base=not args.enable_refiner)
+    # No VAE decoder in base when it outputs latent instead of image.
+    base_info = PipelineInfo(args.version, use_vae=False)
     base = init_pipeline(Txt2ImgXLPipeline, base_info, engine_type, args, max_batch_size, batch_size)
 
-    if args.enable_refiner:
-        refiner_info = PipelineInfo(args.version, is_sd_xl_refiner=True)
-        refiner = init_pipeline(Img2ImgXLPipeline, refiner_info, engine_type, args, max_batch_size, batch_size)
+    refiner_info = PipelineInfo(args.version, is_refiner=True)
+    refiner = init_pipeline(Img2ImgXLPipeline, refiner_info, engine_type, args, max_batch_size, batch_size)
 
-        if engine_type == EngineType.TRT:
-            max_device_memory = max(base.backend.max_device_memory(), refiner.backend.max_device_memory())
-            _, shared_device_memory = cudart.cudaMalloc(max_device_memory)
-            base.backend.activate_engines(shared_device_memory)
-            refiner.backend.activate_engines(shared_device_memory)
+    if engine_type == EngineType.TRT:
+        max_device_memory = max(base.backend.max_device_memory(), refiner.backend.max_device_memory())
+        _, shared_device_memory = cudart.cudaMalloc(max_device_memory)
+        base.backend.activate_engines(shared_device_memory)
+        refiner.backend.activate_engines(shared_device_memory)
 
-        base.load_resources(image_height, image_width, batch_size)
-        refiner.load_resources(image_height, image_width, batch_size)
-    else:
-        if engine_type == EngineType.TRT:
-            max_device_memory = max(base.backend.max_device_memory(), base.backend.max_device_memory())
-            _, shared_device_memory = cudart.cudaMalloc(max_device_memory)
-            base.backend.activate_engines(shared_device_memory)
+    base.load_resources(image_height, image_width, batch_size)
+    refiner.load_resources(image_height, image_width, batch_size)
 
-        base.load_resources(image_height, image_width, batch_size)
-
-    def run_sd_xl_inference(enable_refiner: bool, warmup=False):
+    def run_base_and_refiner(warmup=False):
         images, time_base = base.run(
             prompt,
             negative_prompt,
@@ -86,44 +83,46 @@ if __name__ == "__main__":
             denoising_steps=args.denoising_steps,
             guidance=args.guidance,
             seed=args.seed,
-            return_type="latents" if enable_refiner else "images",
+            return_type="latent",
         )
 
-        if enable_refiner:
-            images, time_refiner = refiner.run(
-                prompt,
-                negative_prompt,
-                images,
-                image_height,
-                image_width,
-                warmup=warmup,
-                denoising_steps=args.denoising_steps,
-                guidance=args.guidance,
-                seed=args.seed,
-            )
-            return images, time_base + time_refiner
-        else:
-            return images, time_base
+        images, time_refiner = refiner.run(
+            prompt,
+            negative_prompt,
+            images,
+            image_height,
+            image_width,
+            warmup=warmup,
+            denoising_steps=args.denoising_steps,
+            guidance=args.guidance,
+            seed=args.seed,
+        )
+
+        return images, time_base + time_refiner
 
     if not args.disable_cuda_graph:
         # inference once to get cuda graph
-        images, _ = run_sd_xl_inference(args.enable_refiner, warmup=True)
+        _, _ = run_base_and_refiner(warmup=True)
 
     print("[I] Warming up ..")
     for _ in range(args.num_warmup_runs):
-        images, _ = run_sd_xl_inference(args.enable_refiner, warmup=True)
+        _, _ = run_base_and_refiner(warmup=True)
 
     print("[I] Running StableDiffusion XL pipeline")
     if args.nvtx_profile:
         cudart.cudaProfilerStart()
-    images, pipeline_time = run_sd_xl_inference(args.enable_refiner, warmup=False)
+    _, latency = run_base_and_refiner(warmup=False)
     if args.nvtx_profile:
         cudart.cudaProfilerStop()
 
     base.teardown()
 
-    if args.enable_refiner:
-        print("|------------|--------------|")
-        print("| {:^10} | {:>9.2f} ms |".format("e2e", pipeline_time))
-        print("|------------|--------------|")
-        refiner.teardown()
+    print("|------------|--------------|")
+    print("| {:^10} | {:>9.2f} ms |".format("e2e", latency))
+    print("|------------|--------------|")
+    refiner.teardown()
+
+
+if __name__ == "__main__":
+    coloredlogs.install(fmt="%(funcName)20s: %(message)s")
+    run_demo()
diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_utils.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_utils.py
index 796e83f70d..3996c8c325 100644
--- a/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_utils.py
+++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_utils.py
@@ -34,7 +34,7 @@ class RawTextArgumentDefaultsHelpFormatter(argparse.ArgumentDefaultsHelpFormatte
 def parse_arguments(is_xl: bool, description: str):
     parser = argparse.ArgumentParser(description=description, formatter_class=RawTextArgumentDefaultsHelpFormatter)
 
-    engines = ["ORT_TRT", "TRT"] if is_xl else ["ORT_CUDA", "ORT_TRT", "TRT"]
+    engines = ["ORT_CUDA", "ORT_TRT", "TRT"]
 
     parser.add_argument(
         "--engine",
@@ -95,7 +95,7 @@ def parse_arguments(is_xl: bool, description: str):
         "--denoising-steps",
         type=int,
         default=30 if is_xl else 50,
-        help="Number of denoising steps" + (" in each of base and refiner." if is_xl else "."),
+        help="Number of denoising steps" + (" in base." if is_xl else "."),
     )
 
     parser.add_argument(
@@ -158,12 +158,6 @@ def parse_arguments(is_xl: bool, description: str):
         "--build-all-tactics", action="store_true", help="Build TensorRT engines using all tactic sources."
     )
 
-    # Pipeline options
-    if is_xl:
-        parser.add_argument(
-            "--enable-refiner", action="store_true", help="Enable refiner and run both base and refiner pipelines."
-        )
-
     args = parser.parse_args()
 
     if (
@@ -203,6 +197,7 @@ def repeat_prompt(args):
         raise ValueError(
             f"`--negative-prompt` must be of type `str` or `str` list, but is {type(args.negative_prompt)}"
         )
+
     if len(args.negative_prompt) == 1:
         negative_prompt = args.negative_prompt * len(prompt)
     else:
@@ -236,16 +231,11 @@ def init_pipeline(pipeline_class, pipeline_info, engine_type, args, max_batch_si
             engine_dir=engine_dir,
             framework_model_dir=framework_model_dir,
             onnx_dir=onnx_dir,
-            onnx_opset=args.onnx_opset,
             opt_image_height=args.height,
             opt_image_width=args.height,
             opt_batch_size=batch_size,
             force_engine_rebuild=args.force_engine_build,
             device_id=torch.cuda.current_device(),
-            disable_cuda_graph_models=[
-                "clip2",  # TODO: Add ArgMax cuda kernel to enable cuda graph for clip2.
-                "unetxl",
-            ],
         )
     elif engine_type == EngineType.ORT_TRT:
         # Build TensorRT EP engines and load pytorch modules
diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/diffusion_models.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/diffusion_models.py
index 7726abb9f9..dc777e2693 100644
--- a/onnxruntime/python/tools/transformers/models/stable_diffusion/diffusion_models.py
+++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/diffusion_models.py
@@ -24,7 +24,7 @@
 import logging
 import os
 import tempfile
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 import onnx
 import onnx_graphsurgeon as gs
@@ -82,43 +82,41 @@ class TrtOptimizer:
 
 
 class PipelineInfo:
-    def __init__(
-        self, version: str, is_inpaint: bool = False, is_sd_xl_refiner: bool = False, use_vae_in_xl_base=False
-    ):
+    def __init__(self, version: str, is_inpaint: bool = False, is_refiner: bool = False, use_vae=False):
         self.version = version
         self._is_inpaint = is_inpaint
-        self._is_sd_xl_refiner = is_sd_xl_refiner
-        self._use_vae_in_xl_base = use_vae_in_xl_base
+        self._is_refiner = is_refiner
+        self._use_vae = use_vae
 
-        if is_sd_xl_refiner:
-            assert self.is_sd_xl()
+        if is_refiner:
+            assert self.is_xl()
 
     def is_inpaint(self) -> bool:
         return self._is_inpaint
 
-    def is_sd_xl(self) -> bool:
+    def is_xl(self) -> bool:
         return "xl" in self.version
 
-    def is_sd_xl_base(self) -> bool:
-        return self.is_sd_xl() and not self._is_sd_xl_refiner
+    def is_xl_base(self) -> bool:
+        return self.is_xl() and not self._is_refiner
 
-    def is_sd_xl_refiner(self) -> bool:
-        return self.is_sd_xl() and self._is_sd_xl_refiner
+    def is_xl_refiner(self) -> bool:
+        return self.is_xl() and self._is_refiner
 
     def use_safetensors(self) -> bool:
-        return self.is_sd_xl()
+        return self.is_xl()
 
     def stages(self) -> List[str]:
-        if self.is_sd_xl_base():
-            return ["clip", "clip2", "unetxl"] + (["vae"] if self._use_vae_in_xl_base else [])
+        if self.is_xl_base():
+            return ["clip", "clip2", "unetxl"] + (["vae"] if self._use_vae else [])
 
-        if self.is_sd_xl_refiner():
+        if self.is_xl_refiner():
             return ["clip2", "unetxl", "vae"]
 
         return ["clip", "unet", "vae"]
 
     def vae_scaling_factor(self) -> float:
-        return 0.13025 if self.is_sd_xl() else 0.18215
+        return 0.13025 if self.is_xl() else 0.18215
 
     @staticmethod
     def supported_versions(is_xl: bool):
@@ -150,7 +148,7 @@ class PipelineInfo:
         elif self.version == "2.1-base":
             return "stabilityai/stable-diffusion-2-1-base"
         elif self.version == "xl-1.0":
-            if self.is_sd_xl_refiner():
+            if self.is_xl_refiner():
                 return "stabilityai/stable-diffusion-xl-refiner-1.0"
             else:
                 return "stabilityai/stable-diffusion-xl-base-1.0"
@@ -166,7 +164,7 @@ class PipelineInfo:
             return 768
         elif self.version in ("2.0", "2.0-base", "2.1", "2.1-base"):
             return 1024
-        elif self.version in ("xl-1.0") and self.is_sd_xl_base():
+        elif self.version in ("xl-1.0") and self.is_xl_base():
             return 768
         else:
             raise ValueError(f"Invalid version {self.version}")
@@ -182,9 +180,9 @@ class PipelineInfo:
             return 768
         elif self.version in ("2.0", "2.0-base", "2.1", "2.1-base"):
             return 1024
-        elif self.version in ("xl-1.0") and self.is_sd_xl_base():
+        elif self.version in ("xl-1.0") and self.is_xl_base():
             return 2048
-        elif self.version in ("xl-1.0") and self.is_sd_xl_refiner():
+        elif self.version in ("xl-1.0") and self.is_xl_refiner():
             return 1280
         else:
             raise ValueError(f"Invalid version {self.version}")
@@ -254,16 +252,16 @@ class BaseModel:
     def load_model(self, framework_model_dir: str, hf_token: str, subfolder: str):
         pass
 
-    def get_input_names(self):
+    def get_input_names(self) -> List[str]:
         pass
 
-    def get_output_names(self):
+    def get_output_names(self) -> List[str]:
         pass
 
-    def get_dynamic_axes(self):
-        return None
+    def get_dynamic_axes(self) -> Dict[str, Dict[int, str]]:
+        pass
 
-    def get_sample_input(self, batch_size, image_height, image_width):
+    def get_sample_input(self, batch_size, image_height, image_width) -> tuple:
         pass
 
     def get_profile_id(self, batch_size, image_height, image_width, static_batch, static_image_shape):
@@ -293,10 +291,10 @@ class BaseModel:
 
     def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_image_shape):
         """For TensorRT"""
-        return None
+        pass
 
     def get_shape_dict(self, batch_size, image_height, image_width):
-        return None
+        pass
 
     def fp32_input_output_names(self) -> List[str]:
         """For CUDA EP, we export ONNX model with FP32 first, then convert it to mixed precision model.
@@ -305,9 +303,16 @@ class BaseModel:
         """
         return []
 
-    def optimize_ort(self, input_onnx_path, optimized_onnx_path, to_fp16=True):
+    def optimize_ort(self, input_onnx_path, optimized_onnx_path, to_fp16=True, fp32_op_list=None, optimize_by_ort=True):
         optimizer = self.get_ort_optimizer()
-        optimizer.optimize(input_onnx_path, optimized_onnx_path, to_fp16, keep_io_types=self.fp32_input_output_names())
+        optimizer.optimize(
+            input_onnx_path,
+            optimized_onnx_path,
+            float16=to_fp16,
+            keep_io_types=self.fp32_input_output_names(),
+            fp32_op_list=fp32_op_list,
+            optimize_by_ort=optimize_by_ort,
+        )
 
     def optimize_trt(self, input_onnx_path, optimized_onnx_path):
         onnx_graph = onnx.load(input_onnx_path)
@@ -382,7 +387,7 @@ class CLIP(BaseModel):
             max_batch_size=max_batch_size,
             embedding_dim=embedding_dim if embedding_dim > 0 else pipeline_info.clip_embedding_dim(),
         )
-        self.output_hidden_state = pipeline_info.is_sd_xl()
+        self.output_hidden_state = pipeline_info.is_xl()
 
         # see https://github.com/huggingface/diffusers/pull/5057 for more information of clip_skip.
         # Clip_skip=1 means that the output of the pre-final layer will be used for computing the prompt embeddings.
@@ -466,11 +471,18 @@ class CLIP(BaseModel):
         onnx_model.add_node(cast_node)
         onnx_model.save_model_to_file(optimized_onnx_path, use_external_data_format=use_external_data_format)
 
-    def optimize_ort(self, input_onnx_path, optimized_onnx_path, to_fp16=True):
+    def optimize_ort(self, input_onnx_path, optimized_onnx_path, to_fp16=True, fp32_op_list=None, optimize_by_ort=True):
         optimizer = self.get_ort_optimizer()
+
         if not self.output_hidden_state:
             optimizer.optimize(
-                input_onnx_path, optimized_onnx_path, to_fp16, keep_io_types=[], keep_outputs=["text_embeddings"]
+                input_onnx_path,
+                optimized_onnx_path,
+                float16=to_fp16,
+                keep_io_types=[],
+                fp32_op_list=fp32_op_list,
+                keep_outputs=["text_embeddings"],
+                optimize_by_ort=optimize_by_ort,
             )
         else:
             with tempfile.TemporaryDirectory() as tmp_dir:
@@ -483,9 +495,11 @@ class CLIP(BaseModel):
                 optimizer.optimize(
                     tmp_model_path,
                     optimized_onnx_path,
-                    to_fp16,
+                    float16=to_fp16,
                     keep_io_types=[],
+                    fp32_op_list=fp32_op_list,
                     keep_outputs=["text_embeddings", "hidden_states"],
+                    optimize_by_ort=optimize_by_ort,
                 )
 
     def optimize_trt(self, input_onnx_path, optimized_onnx_path):
@@ -741,27 +755,47 @@ class UNetXL(BaseModel):
 
 # VAE Decoder
 class VAE(BaseModel):
-    def __init__(self, pipeline_info: PipelineInfo, model, device, max_batch_size):
+    def __init__(
+        self,
+        pipeline_info: PipelineInfo,
+        model,
+        device,
+        max_batch_size,
+        fp16: bool = False,
+        custom_fp16_vae: Optional[str] = None,
+    ):
         super().__init__(
             pipeline_info,
             model=model,
             device=device,
+            fp16=fp16,
             max_batch_size=max_batch_size,
         )
 
+        # For SD XL, need custom trained fp16 model to speed up, and avoid overflow at the same time.
+        self.custom_fp16_vae = custom_fp16_vae
+
     def load_model(self, framework_model_dir, hf_token: Optional[str] = None, subfolder: str = "vae_decoder"):
-        model_dir = os.path.join(framework_model_dir, self.pipeline_info.name(), subfolder)
+        model_name = self.custom_fp16_vae or self.pipeline_info.name()
+
+        model_dir = os.path.join(framework_model_dir, model_name, subfolder)
         if not os.path.exists(model_dir):
-            vae = AutoencoderKL.from_pretrained(
-                self.pipeline_info.name(),
-                subfolder="vae",
-                use_safetensors=self.pipeline_info.use_safetensors(),
-                use_auth_token=hf_token,
-            ).to(self.device)
+            if self.custom_fp16_vae:
+                vae = AutoencoderKL.from_pretrained(self.custom_fp16_vae, torch_dtype=torch.float16).to(self.device)
+            else:
+                vae = AutoencoderKL.from_pretrained(
+                    self.pipeline_info.name(),
+                    subfolder="vae",
+                    use_safetensors=self.pipeline_info.use_safetensors(),
+                    use_auth_token=hf_token,
+                ).to(self.device)
             vae.save_pretrained(model_dir)
         else:
             print(f"Load {self.name} pytorch model from: {model_dir}")
-            vae = AutoencoderKL.from_pretrained(model_dir).to(self.device)
+            if self.custom_fp16_vae:
+                vae = AutoencoderKL.from_pretrained(model_dir, torch_dtype=torch.float16).to(self.device)
+            else:
+                vae = AutoencoderKL.from_pretrained(model_dir).to(self.device)
 
         vae.forward = vae.decode
         return vae
@@ -809,7 +843,7 @@ class VAE(BaseModel):
         return (torch.randn(batch_size, 4, latent_height, latent_width, dtype=torch.float32, device=self.device),)
 
     def fp32_input_output_names(self) -> List[str]:
-        return ["latent", "images"]
+        return [] if self.fp16 else ["latent", "images"]
 
 
 def get_tokenizer(pipeline_info: PipelineInfo, framework_model_dir, hf_token, subfolder="tokenizer"):
@@ -819,7 +853,7 @@ def get_tokenizer(pipeline_info: PipelineInfo, framework_model_dir, hf_token, su
         model = CLIPTokenizer.from_pretrained(
             pipeline_info.name(),
             subfolder=subfolder,
-            use_safetensors=pipeline_info.is_sd_xl(),
+            use_safetensors=pipeline_info.is_xl(),
             use_auth_token=hf_token,
         )
         model.save_pretrained(tokenizer_dir)
diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/engine_builder.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/engine_builder.py
index fdf05ffc79..029125c639 100644
--- a/onnxruntime/python/tools/transformers/models/stable_diffusion/engine_builder.py
+++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/engine_builder.py
@@ -59,7 +59,16 @@ class EngineBuilder:
         self.device = torch.device(device)
         self.torch_device = torch.device(device, torch.cuda.current_device())
         self.stages = pipeline_info.stages()
-        self.vae_torch_fallback = self.pipeline_info.is_sd_xl()
+
+        # TODO: use custom fp16 for ORT_TRT, and no need to fallback to torch.
+        self.vae_torch_fallback = self.pipeline_info.is_xl() and engine_type != EngineType.ORT_CUDA
+
+        # For SD XL, use an VAE that modified to run in fp16 precision without generating NaNs.
+        self.custom_fp16_vae = (
+            "madebyollin/sdxl-vae-fp16-fix"
+            if self.pipeline_info.is_xl() and self.engine_type == EngineType.ORT_CUDA
+            else None
+        )
 
         self.models = {}
         self.engines = {}
@@ -130,7 +139,7 @@ class EngineBuilder:
                 fp16=export_fp16_unet,
                 max_batch_size=self.max_batch_size,
                 unet_dim=4,
-                time_dim=(5 if self.pipeline_info.is_sd_xl_refiner() else 6),
+                time_dim=(5 if self.pipeline_info.is_xl_refiner() else 6),
             )
 
         # VAE Decoder
@@ -140,6 +149,7 @@ class EngineBuilder:
                 None,  # not loaded yet
                 device=self.torch_device,
                 max_batch_size=self.max_batch_size,
+                custom_fp16_vae=self.custom_fp16_vae,
             )
 
             if self.vae_torch_fallback:
@@ -156,8 +166,9 @@ class EngineBuilder:
 
     def vae_decode(self, latents):
         if self.vae_torch_fallback:
-            latents = latents.to(dtype=torch.float32)
-            self.torch_models["vae"] = self.torch_models["vae"].to(dtype=torch.float32)
+            if not self.custom_fp16_vae:
+                latents = latents.to(dtype=torch.float32)
+                self.torch_models["vae"] = self.torch_models["vae"].to(dtype=torch.float32)
             images = self.torch_models["vae"](latents)["sample"]
         else:
             images = self.run_engine("vae", {"latent": latents})["images"]
diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/engine_builder_ort_cuda.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/engine_builder_ort_cuda.py
index 936d04e8a1..11a39b0dec 100644
--- a/onnxruntime/python/tools/transformers/models/stable_diffusion/engine_builder_ort_cuda.py
+++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/engine_builder_ort_cuda.py
@@ -7,24 +7,33 @@ import gc
 import logging
 import os
 import shutil
+from typing import List, Optional
 
 import torch
 from diffusion_models import PipelineInfo
 from engine_builder import EngineBuilder, EngineType
+from ort_utils import CudaSession
 
 import onnxruntime as ort
-from onnxruntime.transformers.io_binding_helper import CudaSession
 
 logger = logging.getLogger(__name__)
 
 
 class OrtCudaEngine(CudaSession):
-    def __init__(self, onnx_path, device_id: int = 0, enable_cuda_graph=False, disable_optimization=False):
+    def __init__(
+        self,
+        onnx_path,
+        device_id: int = 0,
+        enable_cuda_graph: bool = False,
+        disable_optimization: bool = False,
+    ):
         self.onnx_path = onnx_path
         self.provider = "CUDAExecutionProvider"
         self.provider_options = CudaSession.get_cuda_provider_options(device_id, enable_cuda_graph)
+        # self.provider_options["enable_skip_layer_norm_strict_mode"] = True
 
         session_options = ort.SessionOptions()
+
         # When the model has been optimized by onnxruntime, we can disable optimization to save session creation time.
         if disable_optimization:
             session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
@@ -47,6 +56,28 @@ class OrtCudaEngine(CudaSession):
         super().allocate_buffers(shape_dict)
 
 
+class _ModelConfig:
+    """
+    Configuration of one model (like Clip, UNet etc) on ONNX export and optimization for CUDA provider.
+    For example, if you want to use fp32 in layer normalization, set the following:
+        force_fp32_ops=["SkipLayerNormalization", "LayerNormalization"]
+    """
+
+    def __init__(
+        self,
+        onnx_opset_version: int,
+        use_cuda_graph: bool,
+        fp16: bool = True,
+        force_fp32_ops: Optional[List[str]] = None,
+        optimize_by_ort: bool = True,
+    ):
+        self.onnx_opset_version = onnx_opset_version
+        self.use_cuda_graph = use_cuda_graph
+        self.fp16 = fp16
+        self.force_fp32_ops = force_fp32_ops
+        self.optimize_by_ort = optimize_by_ort
+
+
 class OrtCudaEngineBuilder(EngineBuilder):
     def __init__(
         self,
@@ -80,18 +111,59 @@ class OrtCudaEngineBuilder(EngineBuilder):
             use_cuda_graph=use_cuda_graph,
         )
 
+        self.model_config = {}
+
+    def _configure(
+        self,
+        model_name: str,
+        onnx_opset_version: int,
+        use_cuda_graph: bool,
+        fp16: bool = True,
+        force_fp32_ops: Optional[List[str]] = None,
+        optimize_by_ort: bool = True,
+    ):
+        self.model_config[model_name] = _ModelConfig(
+            onnx_opset_version,
+            use_cuda_graph,
+            fp16=fp16,
+            force_fp32_ops=force_fp32_ops,
+            optimize_by_ort=optimize_by_ort,
+        )
+
+    def configure_xl(self, onnx_opset_version: int):
+        self._configure(
+            "clip",
+            onnx_opset_version=onnx_opset_version,
+            use_cuda_graph=self.use_cuda_graph,
+        )
+        self._configure(
+            "clip2",
+            onnx_opset_version=onnx_opset_version,  # TODO: ArgMax-12 is not implemented in CUDA
+            use_cuda_graph=False,  # TODO: fix Runtime Error with cuda graph
+        )
+        self._configure(
+            "unetxl",
+            onnx_opset_version=onnx_opset_version,
+            use_cuda_graph=False,  # TODO: fix Runtime Error with cuda graph
+        )
+
+        self._configure(
+            "vae",
+            onnx_opset_version=onnx_opset_version,
+            use_cuda_graph=self.use_cuda_graph,
+        )
+
     def build_engines(
         self,
-        engine_dir,
-        framework_model_dir,
-        onnx_dir,
-        onnx_opset,
-        opt_image_height=512,
-        opt_image_width=512,
-        opt_batch_size=1,
-        force_engine_rebuild=False,
-        device_id=0,
-        disable_cuda_graph_models=None,
+        engine_dir: str,
+        framework_model_dir: str,
+        onnx_dir: str,
+        onnx_opset_version: int = 17,
+        opt_image_height: int = 512,
+        opt_image_width: int = 512,
+        opt_batch_size: int = 1,
+        force_engine_rebuild: bool = False,
+        device_id: int = 0,
     ):
         self.torch_device = torch.device("cuda", device_id)
         self.load_models(framework_model_dir)
@@ -110,6 +182,13 @@ class OrtCudaEngineBuilder(EngineBuilder):
         if not os.path.isdir(onnx_dir):
             os.makedirs(onnx_dir)
 
+        # Add default configuration if missing
+        if self.pipeline_info.is_xl():
+            self.configure_xl(onnx_opset_version)
+        for model_name in self.models:
+            if model_name not in self.model_config:
+                self.model_config[model_name] = _ModelConfig(onnx_opset_version, self.use_cuda_graph)
+
         # Export models to ONNX
         for model_name, model_obj in self.models.items():
             if model_name == "vae" and self.vae_torch_fallback:
@@ -119,8 +198,12 @@ class OrtCudaEngineBuilder(EngineBuilder):
             onnx_opt_path = self.get_onnx_path(model_name, engine_dir, opt=True)
             if not os.path.exists(onnx_opt_path):
                 if not os.path.exists(onnx_path):
+                    print("----")
                     logger.info("Exporting model: %s", onnx_path)
                     model = model_obj.load_model(framework_model_dir, self.hf_token)
+                    if model_name == "vae":
+                        model.to(torch.float32)
+
                     with torch.inference_mode():
                         # For CUDA EP, export FP32 onnx since some graph fusion only supports fp32 graph pattern.
                         inputs = model_obj.get_sample_input(opt_batch_size, opt_image_height, opt_image_width)
@@ -130,7 +213,7 @@ class OrtCudaEngineBuilder(EngineBuilder):
                             inputs,
                             onnx_path,
                             export_params=True,
-                            opset_version=onnx_opset,
+                            opset_version=self.model_config[model_name].onnx_opset_version,
                             do_constant_folding=True,
                             input_names=model_obj.get_input_names(),
                             output_names=model_obj.get_output_names(),
@@ -144,8 +227,16 @@ class OrtCudaEngineBuilder(EngineBuilder):
 
                 # Run graph optimization and convert to mixed precision (computation in FP16)
                 if not os.path.exists(onnx_opt_path):
+                    print("------")
                     logger.info("Generating optimized model: %s", onnx_opt_path)
-                    model_obj.optimize_ort(onnx_path, onnx_opt_path, to_fp16=True)
+
+                    model_obj.optimize_ort(
+                        onnx_path,
+                        onnx_opt_path,
+                        to_fp16=self.model_config[model_name].fp16,
+                        fp32_op_list=self.model_config[model_name].force_fp32_ops,
+                        optimize_by_ort=self.model_config[model_name].optimize_by_ort,
+                    )
                 else:
                     logger.info("Found cached optimized model: %s", onnx_opt_path)
 
@@ -156,11 +247,15 @@ class OrtCudaEngineBuilder(EngineBuilder):
 
             onnx_opt_path = self.get_onnx_path(model_name, engine_dir, opt=True)
 
-            use_cuda_graph = self.use_cuda_graph
-            if self.use_cuda_graph and disable_cuda_graph_models and model_name in disable_cuda_graph_models:
-                use_cuda_graph = False
+            use_cuda_graph = self.model_config[model_name].use_cuda_graph
+
+            engine = OrtCudaEngine(
+                onnx_opt_path,
+                device_id=device_id,
+                enable_cuda_graph=use_cuda_graph,
+                disable_optimization=False,
+            )
 
-            engine = OrtCudaEngine(onnx_opt_path, device_id=device_id, enable_cuda_graph=use_cuda_graph)
             logger.info("%s options for %s: %s", engine.provider, model_name, engine.provider_options)
             built_engines[model_name] = engine
 
diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/engine_builder_ort_trt.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/engine_builder_ort_trt.py
index a6bbd4ee7e..8a39dc2ed6 100644
--- a/onnxruntime/python/tools/transformers/models/stable_diffusion/engine_builder_ort_trt.py
+++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/engine_builder_ort_trt.py
@@ -12,9 +12,9 @@ import torch
 from cuda import cudart
 from diffusion_models import PipelineInfo
 from engine_builder import EngineBuilder, EngineType
+from ort_utils import CudaSession
 
 import onnxruntime as ort
-from onnxruntime.transformers.io_binding_helper import CudaSession
 
 logger = logging.getLogger(__name__)
 
diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/ort_optimizer.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/ort_optimizer.py
index 24570a6ef6..2078f8d1a4 100644
--- a/onnxruntime/python/tools/transformers/models/stable_diffusion/ort_optimizer.py
+++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/ort_optimizer.py
@@ -51,7 +51,16 @@ class OrtStableDiffusionOptimizer:
             model = onnx.load(str(ort_optimized_model_path), load_external_data=True)
             return self.model_type_class_mapping[self.model_type](model)
 
-    def optimize(self, input_fp32_onnx_path, optimized_onnx_path, float16=True, keep_io_types=False, keep_outputs=None):
+    def optimize(
+        self,
+        input_fp32_onnx_path,
+        optimized_onnx_path,
+        float16=True,
+        keep_io_types=False,
+        fp32_op_list=None,
+        keep_outputs=None,
+        optimize_by_ort=True,
+    ):
         """Optimize onnx model using ONNX Runtime transformers optimizer"""
         logger.info(f"Optimize {input_fp32_onnx_path}...")
         fusion_options = FusionOptions(self.model_type)
@@ -76,6 +85,7 @@ class OrtStableDiffusionOptimizer:
             logger.info("Convert to float16 ...")
             m.convert_float_to_float16(
                 keep_io_types=keep_io_types,
+                op_block_list=fp32_op_list,
             )
 
         use_external_data_format = m.model.ByteSize() >= onnx.checker.MAXIMUM_PROTOBUF
@@ -87,7 +97,7 @@ class OrtStableDiffusionOptimizer:
         # to save session creation time. Another benefit is to inspect the final graph for developing purpose.
         from onnxruntime import __version__ as ort_version
 
-        if version.parse(ort_version) >= version.parse("1.16.0") or not use_external_data_format:
+        if optimize_by_ort and (version.parse(ort_version) >= version.parse("1.16.0") or not use_external_data_format):
             m = self.optimize_by_ort(m, use_external_data_format=use_external_data_format)
 
         m.get_operator_statistics()
diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/ort_utils.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/ort_utils.py
index 5c2145845e..0afa13a0f4 100644
--- a/onnxruntime/python/tools/transformers/models/stable_diffusion/ort_utils.py
+++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/ort_utils.py
@@ -7,16 +7,28 @@ import gc
 import logging
 import os
 import shutil
+import sys
 from typing import Union
 
 import torch
 
 import onnxruntime as ort
-from onnxruntime.transformers.io_binding_helper import CudaSession
 
 logger = logging.getLogger(__name__)
 
 
+def add_transformers_dir_to_path():
+    sys.path.append(os.path.dirname(__file__))
+
+    transformers_dir = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+    if transformers_dir not in sys.path:
+        sys.path.append(transformers_dir)
+
+
+add_transformers_dir_to_path()
+from io_binding_helper import CudaSession  # noqa: E402. Walk-around to test locally
+
+
 # -----------------------------------------------------------------------------------------------------
 # Utilities for CUDA EP
 # -----------------------------------------------------------------------------------------------------
diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_img2img_xl.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_img2img_xl.py
index 0e2aeb6174..faa3f8bfaa 100644
--- a/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_img2img_xl.py
+++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_img2img_xl.py
@@ -29,7 +29,7 @@ from pipeline_stable_diffusion import StableDiffusionPipeline
 
 class Img2ImgXLPipeline(StableDiffusionPipeline):
     """
-    Stable Diffusion Img2Img XL pipeline using NVidia TensorRT.
+    Stable Diffusion Img2Img XL pipeline.
     """
 
     def __init__(self, pipeline_info: PipelineInfo, *args, **kwargs):
@@ -40,7 +40,7 @@ class Img2ImgXLPipeline(StableDiffusionPipeline):
             pipeline_info (PipelineInfo):
                 Version and Type of stable diffusion pipeline.
         """
-        assert pipeline_info.is_sd_xl_refiner()
+        assert pipeline_info.is_xl_refiner()
 
         super().__init__(pipeline_info, *args, **kwargs)
 
@@ -73,12 +73,12 @@ class Img2ImgXLPipeline(StableDiffusionPipeline):
         warmup=False,
         return_type="image",
     ):
-        assert len(prompt) == len(negative_prompt)
+        assert negative_prompt is None or len(prompt) == len(negative_prompt)
 
-        # TODO(tianleiwu): Need we use image_height and image_width for the target size here?
-        original_size = (1024, 1024)
+        original_size = (image_height, image_width)
         crops_coords_top_left = (0, 0)
-        target_size = (1024, 1024)
+        target_size = (image_height, image_width)
+
         strength = 0.3
         aesthetic_score = 6.0
         negative_aesthetic_score = 2.5
@@ -94,6 +94,7 @@ class Img2ImgXLPipeline(StableDiffusionPipeline):
 
             # Initialize timesteps
             timesteps, t_start = self.initialize_timesteps(self.denoising_steps, strength)
+
             latent_timestep = timesteps[:1].repeat(batch_size)
 
             # CLIP text encoder 2
@@ -146,10 +147,10 @@ class Img2ImgXLPipeline(StableDiffusionPipeline):
 
         with torch.inference_mode():
             # VAE decode latent
-            if return_type == "latents":
-                images = latents * self.vae_scaling_factor
+            if return_type == "latent":
+                images = latents
             else:
-                images = self.decode_latent(latents)
+                images = self.decode_latent(latents / self.vae_scaling_factor)
 
             torch.cuda.synchronize()
             e2e_toc = time.perf_counter()
@@ -172,7 +173,7 @@ class Img2ImgXLPipeline(StableDiffusionPipeline):
         guidance=5.0,
         seed=None,
         warmup=False,
-        return_type="images",
+        return_type="image",
     ):
         """
         Run the diffusion pipeline.
@@ -197,7 +198,7 @@ class Img2ImgXLPipeline(StableDiffusionPipeline):
             warmup (bool):
                 Indicate if this is a warmup run.
             return_type (str):
-                It can be "latents" or "images".
+                It can be "latent" or "image".
         """
 
         if self.is_backend_tensorrt():
diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_stable_diffusion.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_stable_diffusion.py
index 87443c9904..e28db2b771 100644
--- a/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_stable_diffusion.py
+++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_stable_diffusion.py
@@ -120,7 +120,7 @@ class StableDiffusionPipeline:
 
         self.stages = pipeline_info.stages()
 
-        self.vae_torch_fallback = self.pipeline_info.is_sd_xl()
+        self.vae_torch_fallback = self.pipeline_info.is_xl()
 
         self.use_cuda_graph = use_cuda_graph
 
@@ -129,6 +129,7 @@ class StableDiffusionPipeline:
 
         self.generator = None
         self.denoising_steps = None
+        self.actual_steps = None
 
         # backend engine
         self.engine_type = engine_type
@@ -142,12 +143,12 @@ class StableDiffusionPipeline:
             raise RuntimeError(f"Backend engine type {engine_type.name} is not supported")
 
         # Load text tokenizer
-        if not self.pipeline_info.is_sd_xl_refiner():
+        if not self.pipeline_info.is_xl_refiner():
             self.tokenizer = get_tokenizer(
                 self.pipeline_info, self.framework_model_dir, self.hf_token, subfolder="tokenizer"
             )
 
-        if self.pipeline_info.is_sd_xl():
+        if self.pipeline_info.is_xl():
             self.tokenizer2 = get_tokenizer(
                 self.pipeline_info, self.framework_model_dir, self.hf_token, subfolder="tokenizer_2"
             )
@@ -219,7 +220,14 @@ class StableDiffusionPipeline:
         return tuple(init_images)
 
     def encode_prompt(
-        self, prompt, negative_prompt, encoder="clip", tokenizer=None, pooled_outputs=False, output_hidden_states=False
+        self,
+        prompt,
+        negative_prompt,
+        encoder="clip",
+        tokenizer=None,
+        pooled_outputs=False,
+        output_hidden_states=False,
+        force_zeros_for_empty_prompt=False,
     ):
         if tokenizer is None:
             tokenizer = self.tokenizer
@@ -247,23 +255,32 @@ class StableDiffusionPipeline:
         if output_hidden_states:
             hidden_states = outputs["hidden_states"].clone()
 
-        # Tokenize negative prompt
-        uncond_input_ids = (
-            tokenizer(
-                negative_prompt,
-                padding="max_length",
-                max_length=tokenizer.model_max_length,
-                truncation=True,
-                return_tensors="pt",
-            )
-            .input_ids.type(torch.int32)
-            .to(self.device)
-        )
+        # Note: negative prompt embedding is not needed for SD XL when guidance < 1
 
-        outputs = self.run_engine(encoder, {"input_ids": uncond_input_ids})
-        uncond_embeddings = outputs["text_embeddings"]
-        if output_hidden_states:
-            uncond_hidden_states = outputs["hidden_states"]
+        # For SD XL base, handle force_zeros_for_empty_prompt
+        is_empty_negative_prompt = all([not i for i in negative_prompt])
+        if force_zeros_for_empty_prompt and is_empty_negative_prompt:
+            uncond_embeddings = torch.zeros_like(text_embeddings)
+            if output_hidden_states:
+                uncond_hidden_states = torch.zeros_like(hidden_states)
+        else:
+            # Tokenize negative prompt
+            uncond_input_ids = (
+                tokenizer(
+                    negative_prompt,
+                    padding="max_length",
+                    max_length=tokenizer.model_max_length,
+                    truncation=True,
+                    return_tensors="pt",
+                )
+                .input_ids.type(torch.int32)
+                .to(self.device)
+            )
+
+            outputs = self.run_engine(encoder, {"input_ids": uncond_input_ids})
+            uncond_embeddings = outputs["text_embeddings"]
+            if output_hidden_states:
+                uncond_hidden_states = outputs["hidden_states"]
 
         # Concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes for classifier free guidance
         text_embeddings = torch.cat([uncond_embeddings, text_embeddings]).to(dtype=torch.float16)
@@ -292,21 +309,21 @@ class StableDiffusionPipeline:
         mask=None,
         masked_image_latents=None,
         guidance=7.5,
-        image_guidance=1.5,
         add_kwargs=None,
     ):
-        assert guidance > 1.0, "Guidance has to be > 1.0"
-        assert image_guidance > 1.0, "Image guidance has to be > 1.0"
+        assert guidance > 1.0, "Guidance has to be > 1.0"  # TODO: remove this constraint
 
         cudart.cudaEventRecord(self.events["denoise-start"], 0)
         if not isinstance(timesteps, torch.Tensor):
             timesteps = self.scheduler.timesteps
+
         for step_index, timestep in enumerate(timesteps):
             if self.nvtx_profile:
                 nvtx_latent_scale = nvtx.start_range(message="latent_scale", color="pink")
 
             # Expand the latents if we are doing classifier free guidance
             latent_model_input = torch.cat([latents] * 2)
+
             latent_model_input = self.scheduler.scale_model_input(
                 latent_model_input, step_offset + step_index, timestep
             )
@@ -322,11 +339,11 @@ class StableDiffusionPipeline:
 
             timestep_float = timestep.float() if timestep.dtype != torch.float32 else timestep
 
-            sample_inp = latent_model_input
-            timestep_inp = timestep_float
-            embeddings_inp = text_embeddings
-
-            params = {"sample": sample_inp, "timestep": timestep_inp, "encoder_hidden_states": embeddings_inp}
+            params = {
+                "sample": latent_model_input,
+                "timestep": timestep_float,
+                "encoder_hidden_states": text_embeddings,
+            }
             if add_kwargs:
                 params.update(add_kwargs)
 
@@ -338,7 +355,7 @@ class StableDiffusionPipeline:
             if self.nvtx_profile:
                 nvtx_latent_step = nvtx.start_range(message="latent_step", color="pink")
 
-            # Perform guidance
+            # perform guidance
             noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
             noise_pred = noise_pred_uncond + guidance * (noise_pred_text - noise_pred_uncond)
 
@@ -350,8 +367,11 @@ class StableDiffusionPipeline:
             if self.nvtx_profile:
                 nvtx.end_range(nvtx_latent_step)
 
-        latents = 1.0 / self.vae_scaling_factor * latents
         cudart.cudaEventRecord(self.events["denoise-stop"], 0)
+
+        # The actual number of steps. It might be different from denoising_steps.
+        self.actual_steps = len(timesteps)
+
         return latents
 
     def encode_image(self, init_image):
@@ -394,7 +414,7 @@ class StableDiffusionPipeline:
         )
         print(
             "| {:^10} | {:>9.2f} ms |".format(
-                "UNet x " + str(self.denoising_steps),
+                "UNet x " + str(self.actual_steps),
                 cudart.cudaEventElapsedTime(self.events["denoise-start"], self.events["denoise-stop"])[1],
             )
         )
@@ -403,6 +423,7 @@ class StableDiffusionPipeline:
                 "VAE-Dec", cudart.cudaEventElapsedTime(self.events["vae-start"], self.events["vae-stop"])[1]
             )
         )
+
         print("|------------|--------------|")
         print("| {:^10} | {:>9.2f} ms |".format("Pipeline", (toc - tic) * 1000.0))
         print("|------------|--------------|")
@@ -413,6 +434,7 @@ class StableDiffusionPipeline:
         images = (
             ((images + 1) * 255 / 2).clamp(0, 255).detach().permute(0, 2, 3, 1).round().type(torch.uint8).cpu().numpy()
         )
+
         from PIL import Image
 
         return [Image.fromarray(images[i]) for i in range(images.shape[0])]
diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_txt2img.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_txt2img.py
index 82f73e8b3c..444b6d9a8c 100644
--- a/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_txt2img.py
+++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_txt2img.py
@@ -52,7 +52,7 @@ class Txt2ImgPipeline(StableDiffusionPipeline):
         guidance=7.5,
         seed=None,
         warmup=False,
-        return_type="latents",
+        return_type="latent",
     ):
         assert len(prompt) == len(negative_prompt)
         batch_size = len(prompt)
@@ -100,7 +100,7 @@ class Txt2ImgPipeline(StableDiffusionPipeline):
         guidance=7.5,
         seed=None,
         warmup=False,
-        return_type="images",
+        return_type="image",
     ):
         """
         Run the diffusion pipeline.
@@ -123,7 +123,7 @@ class Txt2ImgPipeline(StableDiffusionPipeline):
             warmup (bool):
                 Indicate if this is a warmup run.
             return_type (str):
-                type of return. The value can be "latents" or "images".
+                type of return. The value can be "latent" or "image".
         """
         if self.is_backend_tensorrt():
             import tensorrt as trt
diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_txt2img_xl.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_txt2img_xl.py
index d8f00ed619..1b3be143e6 100644
--- a/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_txt2img_xl.py
+++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/pipeline_txt2img_xl.py
@@ -40,7 +40,7 @@ class Txt2ImgXLPipeline(StableDiffusionPipeline):
             pipeline_info (PipelineInfo):
                 Version and Type of stable diffusion pipeline.
         """
-        assert pipeline_info.is_sd_xl_base()
+        assert pipeline_info.is_xl_base()
 
         super().__init__(pipeline_info, *args, **kwargs)
 
@@ -59,14 +59,13 @@ class Txt2ImgXLPipeline(StableDiffusionPipeline):
         guidance=5.0,
         seed=None,
         warmup=False,
-        return_type="images",
+        return_type="image",
     ):
         assert len(prompt) == len(negative_prompt)
 
-        # TODO(tianleiwu): Need we use image_height and image_width for the target size here?
-        original_size = (1024, 1024)
+        original_size = (image_height, image_width)
         crops_coords_top_left = (0, 0)
-        target_size = (1024, 1024)
+        target_size = (image_height, image_width)
         batch_size = len(prompt)
 
         self.set_denoising_steps(denoising_steps)
@@ -86,7 +85,12 @@ class Txt2ImgXLPipeline(StableDiffusionPipeline):
 
             # CLIP text encoder
             text_embeddings = self.encode_prompt(
-                prompt, negative_prompt, encoder="clip", tokenizer=self.tokenizer, output_hidden_states=True
+                prompt,
+                negative_prompt,
+                encoder="clip",
+                tokenizer=self.tokenizer,
+                output_hidden_states=True,
+                force_zeros_for_empty_prompt=True,
             )
             # CLIP text encoder 2
             text_embeddings2, pooled_embeddings2 = self.encode_prompt(
@@ -96,6 +100,7 @@ class Txt2ImgXLPipeline(StableDiffusionPipeline):
                 tokenizer=self.tokenizer2,
                 pooled_outputs=True,
                 output_hidden_states=True,
+                force_zeros_for_empty_prompt=True,
             )
 
             # Merged text embeddings
@@ -112,14 +117,18 @@ class Txt2ImgXLPipeline(StableDiffusionPipeline):
 
             # UNet denoiser
             latents = self.denoise_latent(
-                latents, text_embeddings, denoiser="unetxl", guidance=guidance, add_kwargs=add_kwargs
+                latents,
+                text_embeddings,
+                denoiser="unetxl",
+                guidance=guidance,
+                add_kwargs=add_kwargs,
             )
 
             # VAE decode latent
-            if return_type == "latents":
-                images = latents * self.vae_scaling_factor
+            if return_type == "latent":
+                images = latents
             else:
-                images = self.decode_latent(latents)
+                images = self.decode_latent(latents / self.vae_scaling_factor)
 
             torch.cuda.synchronize()
             e2e_toc = time.perf_counter()
@@ -127,7 +136,7 @@ class Txt2ImgXLPipeline(StableDiffusionPipeline):
             if not warmup:
                 print("SD-XL Base Pipeline")
                 self.print_summary(e2e_tic, e2e_toc, batch_size)
-                if return_type == "images":
+                if return_type != "latent":
                     self.save_images(images, "txt2img-xl", prompt)
 
             return images, (e2e_toc - e2e_tic) * 1000.0
@@ -142,7 +151,7 @@ class Txt2ImgXLPipeline(StableDiffusionPipeline):
         guidance=5.0,
         seed=None,
         warmup=False,
-        return_type="images",
+        return_type="image",
     ):
         """
         Run the diffusion pipeline.
@@ -165,7 +174,7 @@ class Txt2ImgXLPipeline(StableDiffusionPipeline):
             warmup (bool):
                 Indicate if this is a warmup run.
             return_type (str):
-                It can be "latents" or "images".
+                It can be "latent" or "image".
         """
 
         if self.is_backend_tensorrt():