Add SDXL Turbo to demo (#18627)

* Add SDXL Turbo to the demo.
* Change default scheduler to EulerA for XL or Turbo since DDIM does not
work well with small steps.

Example to run the model in demo (See README for instructions):
```
python3 demo_txt2img_xl.py --version xl-turbo --height 512 --width 512 --denoising-steps 1 --scheduler UniPC "little cute gremlin sitting on a bed, cinematic"
```

onnxruntime/python/tools/transformers/models/stable_diffusion/README.md

@@ -85,18 +85,26 @@ If you do not provide prompt, the script will generate different image sizes for
 
 ### Generate an image guided by a text prompt using LCM LoRA
 ```
-python3 demo_txt2img_xl.py "Self-portrait oil painting, a beautiful cyborg with golden hair, 8k" --scheduler LCM --lora-weights latent-consistency/lcm-lora-sdxl --denoising-steps 4
+python3 demo_txt2img_xl.py "Self-portrait oil painting, a beautiful cyborg with golden hair, 8k" --scheduler LCM --lora-weights latent-consistency/lcm-lora-sdxl --denoising-steps 4 --disable-refiner
 ```
 #### Generate an image with SDXL LCM model guided by a text prompt
 ```
 python3 demo_txt2img_xl.py --lcm --disable-refiner "an astronaut riding a rainbow unicorn, cinematic, dramatic"
 ```
 
+#### Generate an image with SDXL Turbo model guided by a text prompt
+It is recommended to use LCM or EuerA scheduler to run SDXL Turbo model.
+```
+python3 demo_txt2img_xl.py --version xl-turbo --height 512 --width 512 --denoising-steps 4 --scheduler LCM "little cute gremlin wearing a jacket, cinematic, vivid colors, intricate masterpiece, golden ratio, highly detailed"
+```
+
 #### Generate an image with a text prompt using a control net
+Control Net is supported for 1.5, SD XL and Turbo models in this demo.
+
 ```
 python3 demo_txt2img.py "Stormtrooper's lecture in beautiful lecture hall" --controlnet-type depth --controlnet-scale 1.0
 
-python3 demo_txt2img_xl.py "young Mona Lisa" --controlnet-type canny --controlnet-scale 0.5 --scheduler UniPC --disable-refiner
+python3 demo_txt2img_xl.py --controlnet-type canny --controlnet-scale 0.5 --version xl-turbo --denoising-steps 2 --scheduler LCM --height 768 --width 768 "portrait of young Mona Lisa with mountain, river and forest in the background"
 ```
 
 ## Optimize Stable Diffusion ONNX models for Hugging Face Diffusers or Optimum

onnxruntime/python/tools/transformers/models/stable_diffusion/demo_txt2img_xl.py

@@ -54,8 +54,12 @@ def load_pipelines(args, batch_size):
     # For TensorRT,  performance of engine built with dynamic shape is very sensitive to the range of image size.
     # Here, we reduce the range of image size for TensorRT to trade-off flexibility and performance.
     # This range can cover most frequent shape of landscape (832x1216), portrait (1216x832) or square (1024x1024).
-    min_image_size = 832 if args.engine != "ORT_CUDA" else 512
-    max_image_size = 1216 if args.engine != "ORT_CUDA" else 2048
+    if args.version == "xl-turbo":
+        min_image_size = 512
+        max_image_size = 768 if args.engine != "ORT_CUDA" else 1024
+    else:
+        min_image_size = 832 if args.engine != "ORT_CUDA" else 512
+        max_image_size = 1216 if args.engine != "ORT_CUDA" else 2048
 
     # No VAE decoder in base when it outputs latent instead of image.
     base_info = PipelineInfo(
@@ -239,12 +243,12 @@ def run_dynamic_shape_demo(args):
         "close-up photography of old man standing in the rain at night, in a street lit by lamps, leica 35mm",
     ]
 
-    # refiner, batch size, height, width, scheduler, steps, prompt, seed, guidance, refiner scheduler, refiner steps, refiner strength
+    # batch size, height, width, scheduler, steps, prompt, seed, guidance, refiner scheduler, refiner steps, refiner strength
     configs = [
         (1, 832, 1216, "UniPC", 8, prompts[0], None, 5.0, "UniPC", 10, 0.3),
         (1, 1024, 1024, "DDIM", 24, prompts[1], None, 5.0, "DDIM", 30, 0.3),
-        (1, 1216, 832, "UniPC", 16, prompts[2], None, 5.0, "UniPC", 10, 0.3),
-        (1, 1344, 768, "DDIM", 24, prompts[3], None, 5.0, "UniPC", 20, 0.3),
+        (1, 1216, 832, "EulerA", 16, prompts[2], 1716921396712843, 5.0, "EulerA", 10, 0.3),
+        (1, 1344, 768, "EulerA", 24, prompts[3], 123698071912362, 5.0, "EulerA", 20, 0.3),
         (2, 640, 1536, "UniPC", 16, prompts[4], 4312973633252712, 5.0, "UniPC", 10, 0.3),
         (2, 1152, 896, "DDIM", 24, prompts[5], 1964684802882906, 5.0, "UniPC", 20, 0.3),
     ]

onnxruntime/python/tools/transformers/models/stable_diffusion/demo_utils.py

@@ -61,7 +61,7 @@ def parse_arguments(is_xl: bool, parser):
     parser.add_argument(
         "--version",
         type=str,
-        default=supported_versions[-1] if is_xl else "1.5",
+        default="xl-1.0" if is_xl else "1.5",
         choices=supported_versions,
         help="Version of Stable Diffusion" + (" XL." if is_xl else "."),
     )
@@ -79,8 +79,8 @@ def parse_arguments(is_xl: bool, parser):
     parser.add_argument(
         "--scheduler",
         type=str,
-        default="DDIM",
-        choices=["DDIM", "UniPC", "LCM"] if is_xl else ["DDIM", "EulerA", "UniPC", "LCM"],
+        default="EulerA" if is_xl else "DDIM",
+        choices=["DDIM", "EulerA", "UniPC", "LCM"],
         help="Scheduler for diffusion process" + " of base" if is_xl else "",
     )
 
@@ -132,8 +132,8 @@ def parse_arguments(is_xl: bool, parser):
         parser.add_argument(
             "--refiner-scheduler",
             type=str,
-            default="DDIM",
-            choices=["DDIM", "UniPC"],
+            default="EulerA",
+            choices=["DDIM", "EulerA", "UniPC"],
             help="Scheduler for diffusion process of refiner.",
         )
 
@@ -244,6 +244,20 @@ def parse_arguments(is_xl: bool, parser):
         args.onnx_opset = 14 if args.engine == "ORT_CUDA" else 17
 
     if is_xl:
+        if args.version == "xl-turbo":
+            if args.guidance > 1.0:
+                print("[I] Use --guidance=0.0 for sdxl-turbo.")
+                args.guidance = 0.0
+            if args.lcm:
+                print("[I] sdxl-turbo cannot use with LCM.")
+                args.lcm = False
+            if args.denoising_steps > 8:
+                print("[I] Use --denoising_steps=4 (no more than 8) for sdxl-turbo.")
+                args.denoising_steps = 4
+            if not args.disable_refiner:
+                print("[I] Disable SDXL refiner to run sdxl-turbo.")
+                args.disable_refiner = True
+
         if args.lcm and args.scheduler != "LCM":
             print("[I] Use --scheduler=LCM for base since LCM is used.")
             args.scheduler = "LCM"
@@ -254,8 +268,8 @@ def parse_arguments(is_xl: bool, parser):
 
     if args.scheduler == "LCM":
         if args.guidance > 1.0:
-            print("[I] Use --guidance=1.0 for base since LCM is used.")
-            args.guidance = 1.0
+            print("[I] Use --guidance=0.0 for base since LCM is used.")
+            args.guidance = 0.0
         if args.denoising_steps > 16:
             print("[I] Use --denoising_steps=8 (no more than 16) for base since LCM is used.")
             args.denoising_steps = 8
@@ -519,7 +533,7 @@ def add_controlnet_arguments(parser, is_xl: bool = False):
         nargs="*",
         type=float,
         default=[],
-        help="The outputs of the controlnet are multiplied by `controlnet_scale` before they are added to the residual in the original unet. Default is 0.35 for SDXL, or 1.0 for SD 1.5",
+        help="The outputs of the controlnet are multiplied by `controlnet_scale` before they are added to the residual in the original unet. Default is 0.5 for SDXL, or 1.0 for SD 1.5",
     )
 
 
@@ -628,12 +642,12 @@ def process_controlnet_arguments(args):
     assert isinstance(args.controlnet_type, list)
     assert isinstance(args.controlnet_scale, list)
     assert isinstance(args.controlnet_image, list)
-    if args.version not in ["1.5", "xl-1.0"]:
-        raise ValueError("This demo only supports ControlNet in Stable Diffusion 1.5 or XL.")
+    if args.version not in ["1.5", "xl-1.0", "xl-turbo"]:
+        raise ValueError("This demo only supports ControlNet in Stable Diffusion 1.5, XL or Turbo.")
 
-    is_xl = args.version == "xl-1.0"
+    is_xl = "xl" in args.version
     if is_xl and len(args.controlnet_type) > 1:
-        raise ValueError("This demo only support one ControlNet for Stable Diffusion XL.")
+        raise ValueError("This demo only support one ControlNet for Stable Diffusion XL or Turbo.")
 
     if len(args.controlnet_image) != 0 and len(args.controlnet_image) != len(args.controlnet_scale):
         raise ValueError(

onnxruntime/python/tools/transformers/models/stable_diffusion/diffusion_models.py

@@ -120,17 +120,23 @@ class PipelineInfo:
     def is_xl(self) -> bool:
         return "xl" in self.version
 
+    def is_xl_turbo(self) -> bool:
+        return self.version == "xl-turbo"
+
     def is_xl_base(self) -> bool:
-        return self.is_xl() and not self._is_refiner
+        return self.version == "xl-1.0" and not self._is_refiner
+
+    def is_xl_base_or_turbo(self) -> bool:
+        return self.is_xl_base() or self.is_xl_turbo()
 
     def is_xl_refiner(self) -> bool:
-        return self.is_xl() and self._is_refiner
+        return self.version == "xl-1.0" and self._is_refiner
 
     def use_safetensors(self) -> bool:
         return self.is_xl()
 
     def stages(self) -> List[str]:
-        if self.is_xl_base():
+        if self.is_xl_base_or_turbo():
             return ["clip", "clip2", "unetxl"] + (["vae"] if self._use_vae else [])
 
         if self.is_xl_refiner():
@@ -153,7 +159,7 @@ class PipelineInfo:
 
     @staticmethod
     def supported_versions(is_xl: bool):
-        return ["xl-1.0"] if is_xl else ["1.4", "1.5", "2.0-base", "2.0", "2.1", "2.1-base"]
+        return ["xl-1.0", "xl-turbo"] if is_xl else ["1.4", "1.5", "2.0-base", "2.0", "2.1", "2.1-base"]
 
     def name(self) -> str:
         if self.version == "1.4":
@@ -185,6 +191,8 @@ class PipelineInfo:
                 return "stabilityai/stable-diffusion-xl-refiner-1.0"
             else:
                 return "stabilityai/stable-diffusion-xl-base-1.0"
+        elif self.version == "xl-turbo":
+            return "stabilityai/sdxl-turbo"
 
         raise ValueError(f"Incorrect version {self.version}")
 
@@ -197,13 +205,13 @@ 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_xl_base():
+        elif self.is_xl_base_or_turbo():
             return 768
         else:
             raise ValueError(f"Invalid version {self.version}")
 
     def clipwithproj_embedding_dim(self):
-        if self.version in ("xl-1.0"):
+        if self.is_xl():
             return 1280
         else:
             raise ValueError(f"Invalid version {self.version}")
@@ -213,9 +221,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_xl_base():
+        elif self.is_xl_base_or_turbo():
             return 2048
-        elif self.version in ("xl-1.0") and self.is_xl_refiner():
+        elif self.is_xl_refiner():
             return 1280
         else:
             raise ValueError(f"Invalid version {self.version}")
@@ -227,7 +235,7 @@ class PipelineInfo:
         return self._max_image_size
 
     def default_image_size(self):
-        if self.is_xl():
+        if self.version == "xl-1.0":
             return 1024
         if self.version in ("2.0", "2.1"):
             return 768
@@ -235,7 +243,7 @@ class PipelineInfo:
 
     @staticmethod
     def supported_controlnet(version="1.5"):
-        if version == "xl-1.0":
+        if version in ("xl-1.0", "xl-turbo"):
             return {
                 "canny": "diffusers/controlnet-canny-sdxl-1.0",
                 "depth": "diffusers/controlnet-depth-sdxl-1.0",

onnxruntime/python/tools/transformers/models/stable_diffusion/diffusion_schedulers.py

@@ -38,6 +38,7 @@ class DDIMScheduler:
         set_alpha_to_one: bool = False,
         steps_offset: int = 1,
         prediction_type: str = "epsilon",
+        timestep_spacing: str = "leading",
     ):
         # this schedule is very specific to the latent diffusion model.
         betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
@@ -61,6 +62,7 @@ class DDIMScheduler:
         self.clip_sample = clip_sample
         self.prediction_type = prediction_type
         self.device = device
+        self.timestep_spacing = timestep_spacing
 
     def configure(self):
         variance = np.zeros(self.num_inference_steps, dtype=np.float32)
@@ -88,12 +90,24 @@ class DDIMScheduler:
 
     def set_timesteps(self, num_inference_steps: int):
         self.num_inference_steps = num_inference_steps
-        step_ratio = self.num_train_timesteps // self.num_inference_steps
-        # creates integer timesteps by multiplying by ratio
-        # casting to int to avoid issues when num_inference_step is power of 3
-        timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
+        if self.timestep_spacing == "leading":
+            step_ratio = self.num_train_timesteps // self.num_inference_steps
+            # creates integer timesteps by multiplying by ratio
+            # casting to int to avoid issues when num_inference_step is power of 3
+            timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
+            timesteps += self.steps_offset
+        elif self.timestep_spacing == "trailing":
+            step_ratio = self.num_train_timesteps / self.num_inference_steps
+            # creates integer timesteps by multiplying by ratio
+            # casting to int to avoid issues when num_inference_step is power of 3
+            timesteps = np.round(np.arange(self.num_train_timesteps, 0, -step_ratio)).astype(np.int64)
+            timesteps -= 1
+        else:
+            raise ValueError(
+                f"{self.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."
+            )
+
         self.timesteps = torch.from_numpy(timesteps).to(self.device)
-        self.timesteps += self.steps_offset
 
     def step(
         self,
@@ -199,12 +213,11 @@ class EulerAncestralDiscreteScheduler:
         beta_start: float = 0.0001,
         beta_end: float = 0.02,
         device="cuda",
-        steps_offset=0,
-        prediction_type="epsilon",
+        steps_offset: int = 1,
+        prediction_type: str = "epsilon",
+        timestep_spacing: str = "trailing",  # set default to trailing for SDXL Turbo
     ):
-        # this schedule is very specific to the latent diffusion model.
         betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
-
         alphas = 1.0 - betas
         self.alphas_cumprod = torch.cumprod(alphas, dim=0)
 
@@ -220,16 +233,38 @@ class EulerAncestralDiscreteScheduler:
         timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=float)[::-1].copy()
         self.timesteps = torch.from_numpy(timesteps)
         self.is_scale_input_called = False
+
+        self._step_index = None
+
         self.device = device
         self.num_train_timesteps = num_train_timesteps
         self.steps_offset = steps_offset
         self.prediction_type = prediction_type
+        self.timestep_spacing = timestep_spacing
 
-    def scale_model_input(self, sample: torch.FloatTensor, idx, timestep, *args, **kwargs) -> torch.FloatTensor:
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
+    def _init_step_index(self, timestep):
         if isinstance(timestep, torch.Tensor):
             timestep = timestep.to(self.timesteps.device)
-        step_index = (self.timesteps == timestep).nonzero().item()
-        sigma = self.sigmas[step_index]
+
+        index_candidates = (self.timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        if len(index_candidates) > 1:
+            step_index = index_candidates[1]
+        else:
+            step_index = index_candidates[0]
+
+        self._step_index = step_index.item()
+
+    def scale_model_input(self, sample: torch.FloatTensor, idx, timestep, *args, **kwargs) -> torch.FloatTensor:
+        if self._step_index is None:
+            self._init_step_index(timestep)
+
+        sigma = self.sigmas[self._step_index]
         sample = sample / ((sigma**2 + 1) ** 0.5)
         self.is_scale_input_called = True
         return sample
@@ -237,13 +272,33 @@ class EulerAncestralDiscreteScheduler:
     def set_timesteps(self, num_inference_steps: int):
         self.num_inference_steps = num_inference_steps
 
-        timesteps = np.linspace(0, self.num_train_timesteps - 1, num_inference_steps, dtype=np.float32)[::-1].copy()
+        if self.timestep_spacing == "linspace":
+            timesteps = np.linspace(0, self.num_train_timesteps - 1, num_inference_steps, dtype=np.float32)[::-1].copy()
+        elif self.timestep_spacing == "leading":
+            step_ratio = self.num_train_timesteps // self.num_inference_steps
+            # creates integer timesteps by multiplying by ratio
+            # casting to int to avoid issues when num_inference_step is power of 3
+            timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.float32)
+            timesteps += self.steps_offset
+        elif self.timestep_spacing == "trailing":
+            step_ratio = self.num_train_timesteps / self.num_inference_steps
+            # creates integer timesteps by multiplying by ratio
+            # casting to int to avoid issues when num_inference_step is power of 3
+            timesteps = (np.arange(self.num_train_timesteps, 0, -step_ratio)).round().copy().astype(np.float32)
+            timesteps -= 1
+        else:
+            raise ValueError(
+                f"{self.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."
+            )
+
         sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
         sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
         sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32)
         self.sigmas = torch.from_numpy(sigmas).to(device=self.device)
         self.timesteps = torch.from_numpy(timesteps).to(device=self.device)
 
+        self._step_index = None
+
     def configure(self):
         dts = np.zeros(self.num_inference_steps, dtype=np.float32)
         sigmas_up = np.zeros(self.num_inference_steps, dtype=np.float32)
@@ -270,8 +325,9 @@ class EulerAncestralDiscreteScheduler:
         timestep,
         generator=None,
     ):
-        step_index = (self.timesteps == timestep).nonzero().item()
-        sigma = self.sigmas[step_index]
+        if self._step_index is None:
+            self._init_step_index(timestep)
+        sigma = self.sigmas[self._step_index]
 
         # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
         if self.prediction_type == "epsilon":
@@ -284,12 +340,15 @@ class EulerAncestralDiscreteScheduler:
                 f"prediction_type given as {self.prediction_type} must be one of `epsilon`, or `v_prediction`"
             )
 
-        sigma_up = self.sigmas_up[idx]
+        sigma_from = self.sigmas[self._step_index]
+        sigma_to = self.sigmas[self._step_index + 1]
+        sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5
+        sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5
 
         # 2. Convert to an ODE derivative
         derivative = (sample - pred_original_sample) / sigma
 
-        dt = self.dts[idx]
+        dt = sigma_down - sigma
 
         prev_sample = sample + derivative * dt
 
@@ -298,11 +357,23 @@ class EulerAncestralDiscreteScheduler:
 
         prev_sample = prev_sample + noise * sigma_up
 
+        # upon completion increase step index by one
+        self._step_index += 1
+
         return prev_sample
 
     def add_noise(self, original_samples, noise, idx, timestep=None):
-        step_index = (self.timesteps == timestep).nonzero().item()
-        noisy_samples = original_samples + noise * self.sigmas[step_index]
+        sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
+        schedule_timesteps = self.timesteps.to(original_samples.device)
+        timesteps = timestep.to(original_samples.device)
+
+        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(original_samples.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        noisy_samples = original_samples + noise * sigma
         return noisy_samples
 
 
@@ -322,6 +393,11 @@ class UniPCMultistepScheduler:
         solver_type: str = "bh2",
         lower_order_final: bool = True,
         disable_corrector: Optional[List[int]] = None,
+        use_karras_sigmas: Optional[bool] = False,
+        timestep_spacing: str = "linspace",
+        steps_offset: int = 0,
+        sigma_min=None,
+        sigma_max=None,
     ):
         self.device = device
         self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
@@ -346,6 +422,9 @@ class UniPCMultistepScheduler:
         self.lower_order_nums = 0
         self.disable_corrector = disable_corrector if disable_corrector else []
         self.last_sample = None
+
+        self._step_index = None
+
         self.num_train_timesteps = num_train_timesteps
         self.solver_order = solver_order
         self.prediction_type = prediction_type
@@ -354,21 +433,58 @@ class UniPCMultistepScheduler:
         self.sample_max_value = sample_max_value
         self.solver_type = solver_type
         self.lower_order_final = lower_order_final
+        self.use_karras_sigmas = use_karras_sigmas
+        self.timestep_spacing = timestep_spacing
+        self.steps_offset = steps_offset
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
 
     def set_timesteps(self, num_inference_steps: int):
-        timesteps = (
-            np.linspace(0, self.num_train_timesteps - 1, num_inference_steps + 1)
-            .round()[::-1][:-1]
-            .copy()
-            .astype(np.int64)
-        )
+        if self.timestep_spacing == "linspace":
+            timesteps = (
+                np.linspace(0, self.num_train_timesteps - 1, num_inference_steps + 1)
+                .round()[::-1][:-1]
+                .copy()
+                .astype(np.int64)
+            )
+        elif self.timestep_spacing == "leading":
+            step_ratio = self.num_train_timesteps // (num_inference_steps + 1)
+            # creates integer timesteps by multiplying by ratio
+            # casting to int to avoid issues when num_inference_step is power of 3
+            timesteps = (np.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int64)
+            timesteps += self.steps_offset
+        elif self.timestep_spacing == "trailing":
+            step_ratio = self.num_train_timesteps / num_inference_steps
+            # creates integer timesteps by multiplying by ratio
+            # casting to int to avoid issues when num_inference_step is power of 3
+            timesteps = np.arange(self.num_train_timesteps, 0, -step_ratio).round().copy().astype(np.int64)
+            timesteps -= 1
+        else:
+            raise ValueError(
+                f"{self.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."
+            )
 
-        # when num_inference_steps == num_train_timesteps, we can end up with
-        # duplicates in timesteps.
-        _, unique_indices = np.unique(timesteps, return_index=True)
-        timesteps = timesteps[np.sort(unique_indices)]
+        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
+        if self.use_karras_sigmas:
+            log_sigmas = np.log(sigmas)
+            sigmas = np.flip(sigmas).copy()
+            sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+            timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round()
+            sigmas = np.concatenate([sigmas, sigmas[-1:]]).astype(np.float32)
+        else:
+            sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
+            sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5
+            sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)
 
-        self.timesteps = torch.from_numpy(timesteps).to(self.device)
+        self.sigmas = torch.from_numpy(sigmas)
+        self.timesteps = torch.from_numpy(timesteps).to(device=self.device, dtype=torch.int64)
 
         self.num_inference_steps = len(timesteps)
 
@@ -378,16 +494,19 @@ class UniPCMultistepScheduler:
         self.lower_order_nums = 0
         self.last_sample = None
 
+        # add an index counter for schedulers that allow duplicated timesteps
+        self._step_index = None
+
     # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
     def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor:
         dtype = sample.dtype
-        batch_size, channels, height, width = sample.shape
+        batch_size, channels, *remaining_dims = sample.shape
 
         if dtype not in (torch.float32, torch.float64):
             sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half
 
         # Flatten sample for doing quantile calculation along each image
-        sample = sample.reshape(batch_size, channels * height * width)
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
 
         abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
 
@@ -395,26 +514,89 @@ class UniPCMultistepScheduler:
         s = torch.clamp(
             s, min=1, max=self.sample_max_value
         )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
-
         s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
         sample = torch.clamp(sample, -s, s) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
 
-        sample = sample.reshape(batch_size, channels, height, width)
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
         sample = sample.to(dtype)
 
         return sample
 
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t
+    def _sigma_to_t(self, sigma, log_sigmas):
+        # get log sigma
+        log_sigma = np.log(np.maximum(sigma, 1e-10))
+
+        # get distribution
+        dists = log_sigma - log_sigmas[:, np.newaxis]
+
+        # get sigmas range
+        low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2)
+        high_idx = low_idx + 1
+
+        low = log_sigmas[low_idx]
+        high = log_sigmas[high_idx]
+
+        # interpolate sigmas
+        w = (low - log_sigma) / (low - high)
+        w = np.clip(w, 0, 1)
+
+        # transform interpolation to time range
+        t = (1 - w) * low_idx + w * high_idx
+        t = t.reshape(sigma.shape)
+        return t
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t
+    def _sigma_to_alpha_sigma_t(self, sigma):
+        alpha_t = 1 / ((sigma**2 + 1) ** 0.5)
+        sigma_t = sigma * alpha_t
+
+        return alpha_t, sigma_t
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
+    def _convert_to_karras(self, in_sigmas: torch.FloatTensor, num_inference_steps) -> torch.FloatTensor:
+        """Constructs the noise schedule of Karras et al. (2022)."""
+
+        sigma_min = self.sigma_min
+        sigma_max = self.sigma_max
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        rho = 7.0  # 7.0 is the value used in the paper
+        ramp = np.linspace(0, 1, num_inference_steps)
+        min_inv_rho = sigma_min ** (1 / rho)
+        max_inv_rho = sigma_max ** (1 / rho)
+        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+        return sigmas
+
     def convert_model_output(
-        self, model_output: torch.FloatTensor, timestep: int, sample: torch.FloatTensor
+        self,
+        model_output: torch.FloatTensor,
+        *args,
+        sample: torch.FloatTensor = None,
+        **kwargs,
     ) -> torch.FloatTensor:
+        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError("missing `sample` as a required keyword argument")
+        if timestep is not None:
+            print(
+                "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma = self.sigmas[self.step_index]
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+
         if self.predict_x0:
             if self.prediction_type == "epsilon":
-                alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep]
                 x0_pred = (sample - sigma_t * model_output) / alpha_t
             elif self.prediction_type == "sample":
                 x0_pred = model_output
             elif self.prediction_type == "v_prediction":
-                alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep]
                 x0_pred = alpha_t * sample - sigma_t * model_output
             else:
                 raise ValueError(
@@ -430,11 +612,9 @@ class UniPCMultistepScheduler:
             if self.prediction_type == "epsilon":
                 return model_output
             elif self.prediction_type == "sample":
-                alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep]
                 epsilon = (sample - alpha_t * model_output) / sigma_t
                 return epsilon
             elif self.prediction_type == "v_prediction":
-                alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep]
                 epsilon = alpha_t * model_output + sigma_t * sample
                 return epsilon
             else:
@@ -446,35 +626,55 @@ class UniPCMultistepScheduler:
     def multistep_uni_p_bh_update(
         self,
         model_output: torch.FloatTensor,
-        prev_timestep: int,
-        sample: torch.FloatTensor,
-        order: int,
+        *args,
+        sample: torch.FloatTensor = None,
+        order: Optional[int] = None,
+        **kwargs,
     ) -> torch.FloatTensor:
-        timestep_list = self.timestep_list
+        prev_timestep = args[0] if len(args) > 0 else kwargs.pop("prev_timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError(" missing `sample` as a required keyword argument")
+        if order is None:
+            if len(args) > 2:
+                order = args[2]
+            else:
+                raise ValueError(" missing `order` as a required keyword argument")
+        if prev_timestep is not None:
+            print(
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
         model_output_list = self.model_outputs
 
-        s0, t = self.timestep_list[-1], prev_timestep
+        # s0 = self.timestep_list[-1]
         m0 = model_output_list[-1]
         x = sample
 
-        lambda_t, lambda_s0 = self.lambda_t[t], self.lambda_t[s0]
-        alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0]
-        sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0]
+        sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
 
         h = lambda_t - lambda_s0
+        device = sample.device
 
         rks = []
         d1s = []
         for i in range(1, order):
-            si = timestep_list[-(i + 1)]
+            si = self.step_index - i
             mi = model_output_list[-(i + 1)]
-            lambda_si = self.lambda_t[si]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
             rk = (lambda_si - lambda_s0) / h
             rks.append(rk)
             d1s.append((mi - m0) / rk)
 
         rks.append(1.0)
-        rks = torch.tensor(rks, device=self.device)
+        rks = torch.tensor(rks, device=device)
 
         r = []
         b = []
@@ -499,13 +699,13 @@ class UniPCMultistepScheduler:
             h_phi_k = h_phi_k / hh - 1 / factorial_i
 
         r = torch.stack(r)
-        b = torch.tensor(b, device=self.device)
+        b = torch.tensor(b, device=device)
 
         if len(d1s) > 0:
             d1s = torch.stack(d1s, dim=1)  # (B, K)
             # for order 2, we use a simplified version
             if order == 2:
-                rhos_p = torch.tensor([0.5], dtype=x.dtype, device=self.device)
+                rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
             else:
                 rhos_p = torch.linalg.solve(r[:-1, :-1], b[:-1])
         else:
@@ -514,14 +714,14 @@ class UniPCMultistepScheduler:
         if self.predict_x0:
             x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
             if d1s is not None:
-                pred_res = torch.einsum("k,bkchw->bchw", rhos_p, d1s)
+                pred_res = torch.einsum("k,bkc...->bc...", rhos_p, d1s)
             else:
                 pred_res = 0
             x_t = x_t_ - alpha_t * b_h * pred_res
         else:
             x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
             if d1s is not None:
-                pred_res = torch.einsum("k,bkchw->bchw", rhos_p, d1s)
+                pred_res = torch.einsum("k,bkc...->bc...", rhos_p, d1s)
             else:
                 pred_res = 0
             x_t = x_t_ - sigma_t * b_h * pred_res
@@ -532,38 +732,63 @@ class UniPCMultistepScheduler:
     def multistep_uni_c_bh_update(
         self,
         this_model_output: torch.FloatTensor,
-        this_timestep: int,
-        last_sample: torch.FloatTensor,
-        # this_sample: torch.FloatTensor,
-        order: int,
+        *args,
+        last_sample: torch.FloatTensor = None,
+        this_sample: torch.FloatTensor = None,
+        order: Optional[int] = None,
+        **kwargs,
     ) -> torch.FloatTensor:
-        timestep_list = self.timestep_list
+        this_timestep = args[0] if len(args) > 0 else kwargs.pop("this_timestep", None)
+        if last_sample is None:
+            if len(args) > 1:
+                last_sample = args[1]
+            else:
+                raise ValueError(" missing`last_sample` as a required keyword argument")
+        if this_sample is None:
+            if len(args) > 2:
+                this_sample = args[2]
+            else:
+                raise ValueError(" missing`this_sample` as a required keyword argument")
+        if order is None:
+            if len(args) > 3:
+                order = args[3]
+            else:
+                raise ValueError(" missing`order` as a required keyword argument")
+        if this_timestep is not None:
+            print(
+                "Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
         model_output_list = self.model_outputs
 
-        s0, t = timestep_list[-1], this_timestep
         m0 = model_output_list[-1]
         x = last_sample
         # x_t = this_sample
         model_t = this_model_output
 
-        lambda_t, lambda_s0 = self.lambda_t[t], self.lambda_t[s0]
-        alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0]
-        sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0]
+        sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[self.step_index - 1]
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
 
         h = lambda_t - lambda_s0
+        device = this_sample.device
 
         rks = []
         d1s = []
         for i in range(1, order):
-            si = timestep_list[-(i + 1)]
+            si = self.step_index - (i + 1)
             mi = model_output_list[-(i + 1)]
-            lambda_si = self.lambda_t[si]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
             rk = (lambda_si - lambda_s0) / h
             rks.append(rk)
             d1s.append((mi - m0) / rk)
 
         rks.append(1.0)
-        rks = torch.tensor(rks, device=self.device)
+        rks = torch.tensor(rks, device=device)
 
         r = []
         b = []
@@ -588,7 +813,7 @@ class UniPCMultistepScheduler:
             h_phi_k = h_phi_k / hh - 1 / factorial_i
 
         r = torch.stack(r)
-        b = torch.tensor(b, device=self.device)
+        b = torch.tensor(b, device=device)
 
         if len(d1s) > 0:
             d1s = torch.stack(d1s, dim=1)
@@ -597,14 +822,14 @@ class UniPCMultistepScheduler:
 
         # for order 1, we use a simplified version
         if order == 1:
-            rhos_c = torch.tensor([0.5], dtype=x.dtype, device=self.device)
+            rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
         else:
             rhos_c = torch.linalg.solve(r, b)
 
         if self.predict_x0:
             x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
             if d1s is not None:
-                corr_res = torch.einsum("k,bkchw->bchw", rhos_c[:-1], d1s)
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], d1s)
             else:
                 corr_res = 0
             d1_t = model_t - m0
@@ -612,7 +837,7 @@ class UniPCMultistepScheduler:
         else:
             x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
             if d1s is not None:
-                corr_res = torch.einsum("k,bkchw->bchw", rhos_c[:-1], d1s)
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], d1s)
             else:
                 corr_res = 0
             d1_t = model_t - m0
@@ -620,6 +845,25 @@ class UniPCMultistepScheduler:
         x_t = x_t.to(x.dtype)
         return x_t
 
+    def _init_step_index(self, timestep):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.to(self.timesteps.device)
+
+        index_candidates = (self.timesteps == timestep).nonzero()
+
+        if len(index_candidates) == 0:
+            step_index = len(self.timesteps) - 1
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        elif len(index_candidates) > 1:
+            step_index = index_candidates[1].item()
+        else:
+            step_index = index_candidates[0].item()
+
+        self._step_index = step_index
+
     def step(
         self,
         model_output: torch.FloatTensor,
@@ -632,29 +876,22 @@ class UniPCMultistepScheduler:
                 "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
             )
 
-        if isinstance(timestep, torch.Tensor):
-            timestep = timestep.to(self.device)
-        step_index = (self.timesteps == timestep).nonzero()
-        if len(step_index) == 0:
-            step_index = len(self.timesteps) - 1
-        else:
-            step_index = step_index.item()
+        if self.step_index is None:
+            self._init_step_index(timestep)
 
-        use_corrector = step_index > 0 and step_index - 1 not in self.disable_corrector and self.last_sample is not None
+        use_corrector = (
+            self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None
+        )
 
-        model_output_convert = self.convert_model_output(model_output, timestep, sample)
+        model_output_convert = self.convert_model_output(model_output, sample=sample)
         if use_corrector:
             sample = self.multistep_uni_c_bh_update(
                 this_model_output=model_output_convert,
-                this_timestep=timestep,
                 last_sample=self.last_sample,
-                # this_sample=sample,
+                this_sample=sample,
                 order=self.this_order,
             )
 
-        # now prepare to run the predictor
-        prev_timestep = 0 if step_index == len(self.timesteps) - 1 else self.timesteps[step_index + 1]
-
         for i in range(self.solver_order - 1):
             self.model_outputs[i] = self.model_outputs[i + 1]
             self.timestep_list[i] = self.timestep_list[i + 1]
@@ -663,7 +900,7 @@ class UniPCMultistepScheduler:
         self.timestep_list[-1] = timestep
 
         if self.lower_order_final:
-            this_order = min(self.solver_order, len(self.timesteps) - step_index)
+            this_order = min(self.solver_order, len(self.timesteps) - self.step_index)
         else:
             this_order = self.solver_order
 
@@ -673,7 +910,6 @@ class UniPCMultistepScheduler:
         self.last_sample = sample
         prev_sample = self.multistep_uni_p_bh_update(
             model_output=model_output,  # pass the original non-converted model output, in case solver-p is used
-            prev_timestep=prev_timestep,
             sample=sample,
             order=self.this_order,
         )
@@ -681,6 +917,9 @@ class UniPCMultistepScheduler:
         if self.lower_order_nums < self.solver_order:
             self.lower_order_nums += 1
 
+        # upon completion increase step index by one
+        self._step_index += 1
+
         if not return_dict:
             return (prev_sample,)
 
@@ -689,7 +928,6 @@ class UniPCMultistepScheduler:
     def scale_model_input(self, sample: torch.FloatTensor, *args, **kwargs) -> torch.FloatTensor:
         return sample
 
-    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise
     def add_noise(
         self,
         original_samples: torch.FloatTensor,
@@ -697,21 +935,18 @@ class UniPCMultistepScheduler:
         idx,
         timesteps: torch.IntTensor,
     ) -> torch.FloatTensor:
-        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
-        alphas_cumprod = self.alphas_cumprod.to(device=self.device, dtype=original_samples.dtype)
-        timesteps = timesteps.to(self.device)
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
+        schedule_timesteps = self.timesteps.to(original_samples.device)
+        timesteps = timesteps.to(original_samples.device)
 
-        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
-        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
-        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
-            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(original_samples.shape):
+            sigma = sigma.unsqueeze(-1)
 
-        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
-        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
-        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
-            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
-
-        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+        noisy_samples = alpha_t * original_samples + sigma_t * noise
         return noisy_samples
 
     def configure(self):

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_xl_base()
+        assert pipeline_info.is_xl_base_or_turbo()
 
         super().__init__(pipeline_info, *args, **kwargs)
 

onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt

@@ -1,5 +1,5 @@
-diffusers==0.23.1
-transformers==4.35.1
+diffusers==0.24.0
+transformers==4.35.2
 numpy>=1.24.1
 accelerate
 onnx==1.14.1
@@ -11,7 +11,7 @@ psutil
 sympy
 controlnet_aux
 # The following are for SDXL
-optimum==1.13.1
+optimum==1.14.1
 safetensors
 invisible_watermark
 # newer version of opencv-python migth encounter module 'cv2.dnn' has no attribute 'DictValue' error