diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/README.md b/onnxruntime/python/tools/transformers/models/stable_diffusion/README.md index 3d00c9cd6b..8b6c2a45be 100644 --- a/onnxruntime/python/tools/transformers/models/stable_diffusion/README.md +++ b/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 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 646e3518fa..bf0d7928be 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 @@ -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), ] 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 f0c83fc507..4fe0f58cae 100644 --- a/onnxruntime/python/tools/transformers/models/stable_diffusion/demo_utils.py +++ b/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( 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 c09aff2f51..3c2aa9f829 100644 --- a/onnxruntime/python/tools/transformers/models/stable_diffusion/diffusion_models.py +++ b/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", diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/diffusion_schedulers.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/diffusion_schedulers.py index 6932c8056c..57cb51bbea 100644 --- a/onnxruntime/python/tools/transformers/models/stable_diffusion/diffusion_schedulers.py +++ b/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): 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 d3387ab6db..fa00354942 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_xl_base() + assert pipeline_info.is_xl_base_or_turbo() super().__init__(pipeline_info, *args, **kwargs) diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt b/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt index a04f05f4b2..8865c1505c 100644 --- a/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt +++ b/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