ONNX Stable Diffusion Inpaint WIP

ONNX Stable Diffusion XL Img2Img WIP
ONNX Stable Diffusion XL Refiner WIP
and lots of fixes.

modules/onnx.py

@@ -25,6 +25,22 @@ class OnnxRuntimeModel(OnnxFakeModule, diffusers.OnnxRuntimeModel):
     def named_modules(self): # dummy
         return ()
 
+    def to(self, *args, **kwargs):
+        from modules.onnx_utils import extract_device
+
+        device = extract_device(args, kwargs)
+        if device is not None:
+            from modules.onnx_ep import TORCH_DEVICE_TO_EP
+
+            self.device = device
+            provider = TORCH_DEVICE_TO_EP[device.type] if device.type in TORCH_DEVICE_TO_EP else self.model._providers
+            path = self.model._model_path
+            sess_options = self.model._sess_options
+            del self.model
+            if provider is not None:
+                self.model = OnnxRuntimeModel.load_model(path, provider, sess_options)
+        return self
+
 
 def preprocess_pipeline(p, refiner_enabled: bool):
     from modules import shared, sd_models
@@ -54,6 +70,13 @@ def preprocess_pipeline(p, refiner_enabled: bool):
 
     if hasattr(shared.sd_model, "preprocess"):
         shared.sd_model = shared.sd_model.preprocess(p)
+    if hasattr(shared.sd_refiner, "preprocess"):
+        if shared.opts.onnx_unload_base:
+            sd_models.unload_model_weights(op='model')
+        shared.sd_refiner = shared.sd_refiner.preprocess(p)
+        if shared.opts.onnx_unload_base:
+            sd_models.reload_model_weights(op='model')
+            shared.sd_model = shared.sd_model.preprocess(p)
 
 
 def initialize():

modules/onnx_ep.py

@@ -21,6 +21,12 @@ EP_TO_NAME = {
     ExecutionProvider.ROCm: "gpu-rocm", # test required
     ExecutionProvider.OpenVINO: "gpu-openvino??", # test required
 }
+TORCH_DEVICE_TO_EP = {
+    "cpu": ExecutionProvider.CPU,
+    "cuda": ExecutionProvider.CUDA,
+    "privateuseone": ExecutionProvider.DirectML,
+    "meta": None,
+}
 
 
 def get_default_execution_provider() -> ExecutionProvider:

modules/onnx_pipelines.py

@@ -12,9 +12,7 @@ from abc import ABCMeta
 from typing import Union, Optional, Callable, Type, Tuple, List, Any, Dict
 from diffusers.pipelines.onnx_utils import ORT_TO_NP_TYPE
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
-from diffusers.pipelines.stable_diffusion_xl import StableDiffusionXLPipelineOutput
 from diffusers.image_processor import VaeImageProcessor, PipelineImageInput
-from optimum.pipelines.diffusers.pipeline_stable_diffusion_xl import rescale_noise_cfg
 from installer import log
 from modules import shared
 from modules.paths import sd_configs_path
@@ -22,8 +20,8 @@ from modules.sd_models import CheckpointInfo
 from modules.processing import StableDiffusionProcessing
 from modules.olive import config
 from modules.onnx import OnnxFakeModule, submodels_sd, submodels_sdxl, submodels_sdxl_refiner
-from modules.onnx_utils import check_pipeline_sdxl, check_cache_onnx, load_init_dict, load_submodel, load_submodels, load_pipeline, get_sess_options, patch_kwargs, construct_refiner_pipeline
-from modules.onnx_ep import ExecutionProvider, EP_TO_NAME, get_provider
+from modules.onnx_utils import extract_device, check_pipeline_sdxl, check_cache_onnx, load_init_dict, load_submodel, load_submodels, get_sess_options, patch_kwargs, load_pipeline
+from modules.onnx_ep import ExecutionProvider, EP_TO_NAME, TORCH_DEVICE_TO_EP, get_provider
 
 
 class OnnxPipelineBase(OnnxFakeModule, diffusers.DiffusionPipeline, metaclass=ABCMeta):
@@ -35,6 +33,40 @@ class OnnxPipelineBase(OnnxFakeModule, diffusers.DiffusionPipeline, metaclass=AB
     def __init__(self):
         self.model_type = self.__class__.__name__
 
+    def to(self, *args, **kwargs):
+        if self.__class__ == OnnxRawPipeline:
+            return super().to(*args, **kwargs)
+
+        expected_modules, _ = self._get_signature_keys(self)
+        for name in expected_modules:
+            if not hasattr(self, name):
+                log.warn(f"Pipeline does not have module '{name}'.")
+                continue
+
+            module = getattr(self, name)
+
+            if isinstance(module, optimum.onnxruntime.modeling_diffusion._ORTDiffusionModelPart):
+                device = extract_device(args, kwargs)
+                if device is None:
+                    return self
+                provider = TORCH_DEVICE_TO_EP[device.type] if device.type in TORCH_DEVICE_TO_EP else module.session._providers
+                path = module.session._model_path
+                sess_options = module.session._sess_options
+                del module
+                delattr(self, name)
+                if provider is not None:
+                    setattr(self, name, diffusers.OnnxRuntimeModel.load_model(path, provider, sess_options))
+
+            if not isinstance(module, diffusers.OnnxRuntimeModel):
+                continue
+
+            try:
+                setattr(self, name, module.to(*args, **kwargs))
+                del module
+            except Exception:
+                log.debug(f"Component device/dtype conversion failed: module={name} args={args}, kwargs={kwargs}")
+        return self
+
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, **_):
         return OnnxRawPipeline(
@@ -69,14 +101,12 @@ class OnnxRawPipeline(OnnxPipelineBase):
     scheduler: Any = None # for Img2Img
 
     def __init__(self, constructor: Type[OnnxPipelineBase], path: os.PathLike):
-        self.model_type = constructor.__name__
         self._is_sdxl = check_pipeline_sdxl(constructor)
         self.from_huggingface_cache = shared.opts.diffusers_dir in os.path.abspath(path)
         self.path = path
-        self.original_filename = os.path.basename(path)
+        self.original_filename = os.path.basename(os.path.dirname(os.path.dirname(path))) if self.from_huggingface_cache else os.path.basename(path)
 
         if os.path.isdir(path):
-            self.is_refiner = self._is_sdxl and "Img2Img" in diffusers.DiffusionPipeline.load_config(path)["_class_name"]
             self.init_dict = load_init_dict(constructor, path)
             self.scheduler = load_submodel(self.path, None, "scheduler", self.init_dict["scheduler"])
         else:
@@ -84,21 +114,23 @@ class OnnxRawPipeline(OnnxPipelineBase):
             try:
                 pipeline = cls.from_single_file(path)
                 self.scheduler = pipeline.scheduler
-                if os.path.isdir(shared.opts.onnx_temp_dir):
-                    shutil.rmtree(shared.opts.onnx_temp_dir)
-                os.mkdir(shared.opts.onnx_temp_dir)
-                pipeline.save_pretrained(shared.opts.onnx_temp_dir)
+                path = shared.opts.onnx_temp_dir
+                if os.path.isdir(path):
+                    shutil.rmtree(path)
+                os.mkdir(path)
+                pipeline.save_pretrained(path)
                 del pipeline
-                self.is_refiner = self._is_sdxl and "Img2Img" in diffusers.DiffusionPipeline.load_config(shared.opts.onnx_temp_dir)["_class_name"]
-                self.init_dict = load_init_dict(constructor, shared.opts.onnx_temp_dir)
+                self.init_dict = load_init_dict(constructor, path)
             except Exception:
-                log.error(f'Failed to load pipeline to optimize: is_sdxl={self._is_sdxl}')
+                log.error(f'Failed to load ONNX pipeline: is_sdxl={self._is_sdxl}')
                 log.warn('Model load failed. Please check Diffusers pipeline in Compute Settings.')
                 return
         if "vae" in self.init_dict:
             del self.init_dict["vae"]
 
-        self.constructor = construct_refiner_pipeline if self.is_refiner else constructor
+        self.is_refiner = self._is_sdxl and "Img2Img" not in constructor.__name__ and "Img2Img" in diffusers.DiffusionPipeline.load_config(path)["_class_name"]
+        self.constructor = OnnxStableDiffusionXLImg2ImgPipeline if self.is_refiner else constructor
+        self.model_type = self.constructor.__name__
         self.submodels = (submodels_sdxl_refiner if self.is_refiner else submodels_sdxl) if self._is_sdxl else submodels_sd
 
     def derive_properties(self, pipeline: diffusers.DiffusionPipeline):
@@ -345,8 +377,12 @@ class OnnxRawPipeline(OnnxPipelineBase):
         config.height = p.height
         config.batch_size = p.batch_size
 
-        config.cross_attention_dim = 2048 if self._is_sdxl else (256 + p.height)
-        config.time_ids_size = 6 if self._is_sdxl and not self.is_refiner else 5
+        if self._is_sdxl and not self.is_refiner:
+            config.cross_attention_dim = 2048
+            config.time_ids_size = 6
+        else:
+            config.cross_attention_dim = 256 + p.height
+            config.time_ids_size = 5
 
         if not disable_classifier_free_guidance and "turbo" in str(self.path).lower():
             log.warning("It looks like you are trying to run a Turbo model with CFG Scale, which will lead to 'size mismatch' or 'unexpected parameter' error.")
@@ -357,23 +393,23 @@ class OnnxRawPipeline(OnnxPipelineBase):
 
         converted_dir = self.convert(in_dir)
         if converted_dir is None:
-            log.error('Failed to convert model. The generation will fall back to unconverted one.')
-            return self.derive_properties(load_pipeline(diffusers.StableDiffusionXLPipeline if self._is_sdxl else diffusers.StableDiffusionPipeline, self.path, **kwargs))
+            log.error('Failed to convert model.')
+            return
         out_dir = converted_dir
 
         if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "olive-ai":
             log.warning("Olive implementation is experimental. It contains potentially an issue and is subject to change at any time.")
             if p.width != p.height:
-                log.warning("Olive detected different width and height. The quality of the result is not guaranteed.")
-            if shared.opts.olive_static_dims:
+                log.warning("Olive: different width and height are detected. The quality of the result is not guaranteed.")
+            if shared.opts.olive_static_dims and not self.is_refiner: # Session Conflict
                 kwargs["sess_options"] = get_sess_options(p.batch_size if disable_classifier_free_guidance else p.batch_size * 2, p.height, p.width, self._is_sdxl)
             optimized_dir = self.optimize(converted_dir)
             if optimized_dir is None:
-                log.error('Failed to optimize pipeline. The generation will fall back to unoptimized one.')
-                return self.derive_properties(load_pipeline(diffusers.OnnxStableDiffusionXLPipeline if self._is_sdxl else diffusers.OnnxStableDiffusionPipeline, converted_dir, **kwargs))
+                log.error('Olive: failed to optimize pipeline. The generation will fall back to unoptimized one.')
+                return self.derive_properties(load_pipeline(self.constructor, converted_dir, **kwargs))
             out_dir = optimized_dir
 
-        pipeline = self.derive_properties(load_pipeline(diffusers.OnnxStableDiffusionXLPipeline if self._is_sdxl else diffusers.OnnxStableDiffusionPipeline, out_dir, **kwargs))
+        pipeline = self.derive_properties(load_pipeline(self.constructor, out_dir, **kwargs))
 
         if not shared.opts.onnx_cache_converted:
             shutil.rmtree(converted_dir)
@@ -382,6 +418,22 @@ class OnnxRawPipeline(OnnxPipelineBase):
         return pipeline
 
 
+def extract_generator_seed(generator: Union[torch.Generator, List[torch.Generator]]) -> List[int]:
+    if isinstance(generator, list):
+        generator = [g.seed() for g in generator]
+    else:
+        generator = [generator.seed()]
+    return generator
+
+
+def randn_tensor(shape, dtype, generator: Union[torch.Generator, List[torch.Generator], int, List[int]]):
+    if hasattr(generator, "seed") or (isinstance(generator, list) and hasattr(generator[0], "seed")):
+        generator = extract_generator_seed(generator)
+        if len(generator) == 1:
+            generator = generator[0]
+    return np.random.default_rng(generator).standard_normal(shape).astype(dtype)
+
+
 def prepare_latents(
     init_noise_sigma: float,
     batch_size: int,
@@ -401,12 +453,7 @@ def prepare_latents(
         )
 
     if latents is None:
-        if isinstance(generator, list):
-            generator = [g.seed() for g in generator]
-            if len(generator) == 1:
-                generator = generator[0]
-
-        latents = np.random.default_rng(generator).standard_normal(shape).astype(dtype)
+        latents = randn_tensor(shape, dtype, generator)
     elif latents.shape != shape:
         raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
 
@@ -425,11 +472,11 @@ class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline, OnnxPip
         vae_encoder: diffusers.OnnxRuntimeModel,
         vae_decoder: diffusers.OnnxRuntimeModel,
         text_encoder: diffusers.OnnxRuntimeModel,
-        tokenizer,
+        tokenizer: Any,
         unet: diffusers.OnnxRuntimeModel,
-        scheduler,
+        scheduler: Any,
         safety_checker: diffusers.OnnxRuntimeModel,
-        feature_extractor,
+        feature_extractor: Any,
         requires_safety_checker: bool = True
     ):
         super().__init__(vae_encoder, vae_decoder, text_encoder, tokenizer, unet, scheduler, safety_checker, feature_extractor, requires_safety_checker)
@@ -535,7 +582,7 @@ class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline, OnnxPip
 
             # call the callback, if provided
             if callback is not None and i % callback_steps == 0:
-                callback(i, t, torch.from_numpy(latents))
+                callback(i, t, latents)
 
         latents /= self.vae_decoder.config.get("scaling_factor", 0.18215)
 
@@ -587,11 +634,11 @@ class OnnxStableDiffusionImg2ImgPipeline(diffusers.OnnxStableDiffusionImg2ImgPip
         vae_encoder: diffusers.OnnxRuntimeModel,
         vae_decoder: diffusers.OnnxRuntimeModel,
         text_encoder: diffusers.OnnxRuntimeModel,
-        tokenizer,
+        tokenizer: Any,
         unet: diffusers.OnnxRuntimeModel,
-        scheduler,
+        scheduler: Any,
         safety_checker: diffusers.OnnxRuntimeModel,
-        feature_extractor,
+        feature_extractor: Any,
         requires_safety_checker: bool = True
     ):
         super().__init__(vae_encoder, vae_decoder, text_encoder, tokenizer, unet, scheduler, safety_checker, feature_extractor, requires_safety_checker)
@@ -672,13 +719,8 @@ class OnnxStableDiffusionImg2ImgPipeline(diffusers.OnnxStableDiffusionImg2ImgPip
         timesteps = self.scheduler.timesteps.numpy()[-init_timestep]
         timesteps = np.array([timesteps] * batch_size * num_images_per_prompt)
 
-        if isinstance(generator, list):
-            generator = [g.seed() for g in generator]
-            if len(generator) == 1:
-                generator = generator[0]
-
         # add noise to latents using the timesteps
-        noise = np.random.default_rng(generator).standard_normal(init_latents.shape).astype(latents_dtype)
+        noise = randn_tensor(init_latents.shape, latents_dtype, generator)
         init_latents = self.scheduler.add_noise(
             torch.from_numpy(init_latents), torch.from_numpy(noise), torch.from_numpy(timesteps)
         )
@@ -728,7 +770,7 @@ class OnnxStableDiffusionImg2ImgPipeline(diffusers.OnnxStableDiffusionImg2ImgPip
 
             # call the callback, if provided
             if callback is not None and i % callback_steps == 0:
-                callback(i, t, torch.from_numpy(latents))
+                callback(i, t, latents)
 
         latents /= scaling_factor
 
@@ -780,15 +822,17 @@ class OnnxStableDiffusionInpaintPipeline(diffusers.OnnxStableDiffusionInpaintPip
         vae_encoder: diffusers.OnnxRuntimeModel,
         vae_decoder: diffusers.OnnxRuntimeModel,
         text_encoder: diffusers.OnnxRuntimeModel,
-        tokenizer,
+        tokenizer: Any,
         unet: diffusers.OnnxRuntimeModel,
-        scheduler,
+        scheduler: Any,
         safety_checker: diffusers.OnnxRuntimeModel,
-        feature_extractor,
+        feature_extractor: Any,
         requires_safety_checker: bool = True
     ):
         super().__init__(vae_encoder, vae_decoder, text_encoder, tokenizer, unet, scheduler, safety_checker, feature_extractor, requires_safety_checker)
 
+        self.vae_scale_factor = 2 ** (len(self.vae_decoder.config.get("block_out_channels")) - 1)
+
     @torch.no_grad()
     def __call__(
         self,
@@ -856,7 +900,7 @@ class OnnxStableDiffusionInpaintPipeline(diffusers.OnnxStableDiffusionInpaintPip
             generator,
             latents,
             num_channels_latents,
-            8,
+            self.vae_scale_factor,
         )
 
         scaling_factor = self.vae_decoder.config.get("scaling_factor", 0.18215)
@@ -940,7 +984,7 @@ class OnnxStableDiffusionInpaintPipeline(diffusers.OnnxStableDiffusionInpaintPip
             # call the callback, if provided
             if callback is not None and i % callback_steps == 0:
                 step_idx = i // getattr(self.scheduler, "order", 1)
-                callback(step_idx, t, torch.from_numpy(latents))
+                callback(step_idx, t, latents)
 
         latents /= scaling_factor
 
@@ -987,191 +1031,24 @@ class OnnxStableDiffusionXLPipeline(OnnxPipelineBase, optimum.onnxruntime.ORTSta
 
     def __init__(
         self,
-        vae_decoder,
-        text_encoder,
-        unet,
+        vae_decoder: ort.InferenceSession,
+        text_encoder: ort.InferenceSession,
+        unet: ort.InferenceSession,
         config: Dict[str, Any],
-        tokenizer,
-        scheduler,
-        feature_extractor = None,
-        vae_encoder = None,
-        text_encoder_2 = None,
-        tokenizer_2 = None,
+        tokenizer: Any,
+        scheduler: Any,
+        feature_extractor: Any = None,
+        vae_encoder: Optional[ort.InferenceSession] = None,
+        text_encoder_2: Optional[ort.InferenceSession] = None,
+        tokenizer_2: Any = None,
         use_io_binding: bool | None = None,
         model_save_dir = None,
         add_watermarker: bool | None = None
     ):
         super(optimum.onnxruntime.ORTStableDiffusionXLPipeline, self).__init__(vae_decoder, text_encoder, unet, config, tokenizer, scheduler, feature_extractor, vae_encoder, text_encoder_2, tokenizer_2, use_io_binding, model_save_dir, add_watermarker)
 
-    # Adapted from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.__call__
-    def __call__(
-        self,
-        prompt: Optional[Union[str, List[str]]] = None,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
-        num_inference_steps: int = 50,
-        guidance_scale: float = 5.0,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        num_images_per_prompt: int = 1,
-        eta: float = 0.0,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[np.ndarray] = None,
-        prompt_embeds: Optional[np.ndarray] = None,
-        negative_prompt_embeds: Optional[np.ndarray] = None,
-        pooled_prompt_embeds: Optional[np.ndarray] = None,
-        negative_pooled_prompt_embeds: Optional[np.ndarray] = None,
-        output_type: str = "pil",
-        return_dict: bool = True,
-        callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
-        callback_steps: int = 1,
-        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
-        guidance_rescale: float = 0.0,
-        original_size: Optional[Tuple[int, int]] = None,
-        crops_coords_top_left: Tuple[int, int] = (0, 0),
-        target_size: Optional[Tuple[int, int]] = None,
-    ):
-        # 0. Default height and width to unet
-        height = height or self.unet.config["sample_size"] * self.vae_scale_factor
-        width = width or self.unet.config["sample_size"] * self.vae_scale_factor
-
-        original_size = original_size or (height, width)
-        target_size = target_size or (height, width)
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(
-            prompt,
-            height,
-            width,
-            callback_steps,
-            negative_prompt,
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        )
-
-        # 2. Define call parameters
-        if isinstance(prompt, str):
-            batch_size = 1
-        elif isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        if generator is None:
-            generator = torch.Generator("cpu")
-
-        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
-        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
-        # corresponds to doing no classifier free guidance.
-        do_classifier_free_guidance = guidance_scale > 1.0
-
-        # 3. Encode input prompt
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        ) = self._encode_prompt(
-            prompt,
-            num_images_per_prompt,
-            do_classifier_free_guidance,
-            negative_prompt,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
-        )
-
-        # 4. Prepare timesteps
-        self.scheduler.set_timesteps(num_inference_steps)
-        timesteps = self.scheduler.timesteps
-
-        # 5. Prepare latent variables
-        latents = prepare_latents(
-            self.scheduler.init_noise_sigma,
-            batch_size * num_images_per_prompt,
-            height,
-            width,
-            prompt_embeds.dtype,
-            generator,
-            latents,
-            self.unet.config.get("in_channels", 4),
-            self.vae_scale_factor,
-        )
-
-        # 6. Prepare extra step kwargs
-        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-
-        # 7. Prepare added time ids & embeddings
-        add_text_embeds = pooled_prompt_embeds
-        add_time_ids = (original_size + crops_coords_top_left + target_size,)
-        add_time_ids = np.array(add_time_ids, dtype=prompt_embeds.dtype)
-
-        if do_classifier_free_guidance:
-            prompt_embeds = np.concatenate((negative_prompt_embeds, prompt_embeds), axis=0)
-            add_text_embeds = np.concatenate((negative_pooled_prompt_embeds, add_text_embeds), axis=0)
-            add_time_ids = np.concatenate((add_time_ids, add_time_ids), axis=0)
-        add_time_ids = np.repeat(add_time_ids, batch_size * num_images_per_prompt, axis=0)
-
-        # Adapted from diffusers to extend it for other runtimes than ORT
-        timestep_dtype = self.unet.input_dtype.get("timestep", np.float32)
-
-        # 8. Denoising loop
-        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
-        for i, t in enumerate(self.progress_bar(timesteps)):
-            # expand the latents if we are doing classifier free guidance
-            latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents
-            latent_model_input = self.scheduler.scale_model_input(torch.from_numpy(latent_model_input), t)
-            latent_model_input = latent_model_input.cpu().numpy()
-
-            # predict the noise residual
-            timestep = np.array([t], dtype=timestep_dtype)
-            noise_pred = self.unet(
-                sample=latent_model_input,
-                timestep=timestep,
-                encoder_hidden_states=prompt_embeds,
-                text_embeds=add_text_embeds,
-                time_ids=add_time_ids,
-            )
-            noise_pred = noise_pred[0]
-
-            # perform guidance
-            if do_classifier_free_guidance:
-                noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2)
-                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-                if guidance_rescale > 0.0:
-                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
-                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
-
-            # compute the previous noisy sample x_t -> x_t-1
-            scheduler_output = self.scheduler.step(
-                torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs
-            )
-            latents = scheduler_output.prev_sample.numpy()
-
-            # call the callback, if provided
-            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                if callback is not None and i % callback_steps == 0:
-                    callback(i, t, torch.from_numpy(latents))
-
-        if output_type == "latent":
-            image = latents
-        else:
-            latents /= self.vae_decoder.config.get("scaling_factor", 0.18215)
-            # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1
-            image = np.concatenate(
-                [self.vae_decoder(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])]
-            )
-            # apply watermark if available
-            if self.watermark is not None:
-                image = self.watermark.apply_watermark(image)
-            image = self.image_processor.postprocess(image, output_type=output_type)
-
-        if not return_dict:
-            return (image,)
-
-        return StableDiffusionXLPipelineOutput(images=image)
+    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, generator, latents=None):
+        return prepare_latents(self.scheduler.init_noise_sigma, batch_size, height, width, dtype, generator, latents, num_channels_latents, self.vae_scale_factor)
 
 
 class OnnxStableDiffusionXLImg2ImgPipeline(OnnxPipelineBase, optimum.onnxruntime.ORTStableDiffusionXLImg2ImgPipeline):
@@ -1180,186 +1057,44 @@ class OnnxStableDiffusionXLImg2ImgPipeline(OnnxPipelineBase, optimum.onnxruntime
 
     def __init__(
         self,
-        vae_decoder,
-        text_encoder,
-        unet,
+        vae_decoder: ort.InferenceSession,
+        text_encoder: ort.InferenceSession,
+        unet: ort.InferenceSession,
         config: Dict[str, Any],
-        tokenizer,
-        scheduler,
+        tokenizer: Any,
+        scheduler: Any,
         feature_extractor = None,
-        vae_encoder = None,
-        text_encoder_2 = None,
-        tokenizer_2 = None,
+        vae_encoder: Optional[ort.InferenceSession] = None,
+        text_encoder_2: Optional[ort.InferenceSession] = None,
+        tokenizer_2: Any = None,
         use_io_binding: bool | None = None,
         model_save_dir = None,
         add_watermarker: bool | None = None
     ):
         super(optimum.onnxruntime.ORTStableDiffusionXLImg2ImgPipeline, self).__init__(vae_decoder, text_encoder, unet, config, tokenizer, scheduler, feature_extractor, vae_encoder, text_encoder_2, tokenizer_2, use_io_binding, model_save_dir, add_watermarker)
 
-    # Adapted from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.__call__
-    def __call__(
-        self,
-        prompt: Optional[Union[str, List[str]]] = None,
-        image: PipelineImageInput = None,
-        strength: float = 0.3,
-        num_inference_steps: int = 50,
-        guidance_scale: float = 5.0,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        num_images_per_prompt: int = 1,
-        eta: float = 0.0,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[np.ndarray] = None,
-        prompt_embeds: Optional[np.ndarray] = None,
-        negative_prompt_embeds: Optional[np.ndarray] = None,
-        pooled_prompt_embeds: Optional[np.ndarray] = None,
-        negative_pooled_prompt_embeds: Optional[np.ndarray] = None,
-        output_type: str = "pil",
-        return_dict: bool = True,
-        callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
-        callback_steps: int = 1,
-        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
-        guidance_rescale: float = 0.0,
-        original_size: Optional[Tuple[int, int]] = None,
-        crops_coords_top_left: Tuple[int, int] = (0, 0),
-        target_size: Optional[Tuple[int, int]] = None,
-        aesthetic_score: float = 6.0,
-        negative_aesthetic_score: float = 2.5,
-    ):
-        # 0. Check inputs. Raise error if not correct
-        self.check_inputs(prompt, strength, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds)
+    def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, generator=None):
+        batch_size = batch_size * num_images_per_prompt
 
-        # 1. Define call parameters
-        if isinstance(prompt, str):
-            batch_size = 1
-        elif isinstance(prompt, list):
-            batch_size = len(prompt)
+        if image.shape[1] == 4:
+            init_latents = image
         else:
-            batch_size = prompt_embeds.shape[0]
+            init_latents = self.vae_encoder(sample=image)[0] * self.vae_decoder.config.get("scaling_factor", 0.18215)
 
-        if generator is None:
-            generator = torch.Generator("cpu")
-
-        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
-        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
-        # corresponds to doing no classifier free guidance.
-        do_classifier_free_guidance = guidance_scale > 1.0
-
-        # 2. Encode input prompt
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        ) = self._encode_prompt(
-            prompt,
-            num_images_per_prompt,
-            do_classifier_free_guidance,
-            negative_prompt,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
-        )
-
-        # 3. Preprocess image
-        image = self.image_processor.preprocess(image)
-
-        # 4. Prepare timesteps
-        self.scheduler.set_timesteps(num_inference_steps)
-
-        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength)
-        latent_timestep = np.repeat(timesteps[:1], batch_size * num_images_per_prompt, axis=0)
-        timestep_dtype = self.unet.input_dtype.get("timestep", np.float32)
-
-        latents_dtype = prompt_embeds.dtype
-        image = image.astype(latents_dtype)
-
-        # 5. Prepare latent variables
-        latents = self.prepare_latents(
-            image, latent_timestep, batch_size, num_images_per_prompt, latents_dtype, generator
-        )
-
-        # 6. Prepare extra step kwargs
-        extra_step_kwargs = {}
-        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        if accepts_eta:
-            extra_step_kwargs["eta"] = eta
-
-        height, width = latents.shape[-2:]
-        height = height * self.vae_scale_factor
-        width = width * self.vae_scale_factor
-        original_size = original_size or (height, width)
-        target_size = target_size or (height, width)
-
-        # 8. Prepare added time ids & embeddings
-        add_text_embeds = pooled_prompt_embeds
-        add_time_ids, add_neg_time_ids = self._get_add_time_ids(
-            original_size,
-            crops_coords_top_left,
-            target_size,
-            aesthetic_score,
-            negative_aesthetic_score,
-            dtype=prompt_embeds.dtype,
-        )
-
-        if do_classifier_free_guidance:
-            prompt_embeds = np.concatenate((negative_prompt_embeds, prompt_embeds), axis=0)
-            add_text_embeds = np.concatenate((negative_pooled_prompt_embeds, add_text_embeds), axis=0)
-            add_time_ids = np.concatenate((add_time_ids, add_time_ids), axis=0)
-        add_time_ids = np.repeat(add_time_ids, batch_size * num_images_per_prompt, axis=0)
-
-        # 8. Denoising loop
-        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
-        for i, t in enumerate(self.progress_bar(timesteps)):
-            # expand the latents if we are doing classifier free guidance
-            latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents
-            latent_model_input = self.scheduler.scale_model_input(torch.from_numpy(latent_model_input), t)
-            latent_model_input = latent_model_input.cpu().numpy()
-
-            # predict the noise residual
-            timestep = np.array([t], dtype=timestep_dtype)
-            noise_pred = self.unet(
-                sample=latent_model_input,
-                timestep=timestep,
-                encoder_hidden_states=prompt_embeds,
-                text_embeds=add_text_embeds,
-                time_ids=add_time_ids,
+        if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:
+            # expand init_latents for batch_size
+            additional_image_per_prompt = batch_size // init_latents.shape[0]
+            init_latents = np.concatenate([init_latents] * additional_image_per_prompt, axis=0)
+        elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0:
+            raise ValueError(
+                f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
             )
-            noise_pred = noise_pred[0]
-
-            # perform guidance
-            if do_classifier_free_guidance:
-                noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2)
-                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-                if guidance_rescale > 0.0:
-                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
-                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
-
-            # compute the previous noisy sample x_t -> x_t-1
-            scheduler_output = self.scheduler.step(
-                torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs
-            )
-            latents = scheduler_output.prev_sample.numpy()
-
-            # call the callback, if provided
-            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                if callback is not None and i % callback_steps == 0:
-                    callback(i, t, torch.from_numpy(latents))
-
-        if output_type == "latent":
-            image = latents
         else:
-            latents /= self.vae_decoder.config.get("scaling_factor", 0.18215)
-            # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1
-            image = np.concatenate(
-                [self.vae_decoder(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])]
-            )
-            # apply watermark if available
-            if self.watermark is not None:
-                image = self.watermark.apply_watermark(image)
-            image = self.image_processor.postprocess(image, output_type=output_type)
+            init_latents = np.concatenate([init_latents], axis=0)
 
-        if not return_dict:
-            return (image,)
-
-        return StableDiffusionXLPipelineOutput(images=image)
+        # add noise to latents using the timesteps
+        noise = randn_tensor(init_latents.shape, dtype, generator)
+        init_latents = self.scheduler.add_noise(
+            torch.from_numpy(init_latents), torch.from_numpy(noise), torch.from_numpy(timestep)
+        )
+        return init_latents.numpy()

modules/onnx_utils.py

@@ -1,9 +1,9 @@
 import os
 import importlib
 from typing import Type, Tuple, Union, List, Dict, Any
+import torch
 import diffusers
 import onnxruntime as ort
-from installer import log
 
 
 def get_sess_options(hidden_batch_size: int, height: int, width: int, is_sdxl: bool) -> ort.SessionOptions:
@@ -24,6 +24,15 @@ def get_sess_options(hidden_batch_size: int, height: int, width: int, is_sdxl: b
     return sess_options
 
 
+def extract_device(args: List, kwargs: Dict):
+    device = kwargs.get("device", None)
+    if device is None:
+        for arg in args:
+            if isinstance(arg, torch.device):
+                device = arg
+    return device
+
+
 def load_init_dict(cls: Type[diffusers.DiffusionPipeline], path: os.PathLike):
     merged: Dict[str, Any] = {}
     extracted = cls.extract_init_dict(diffusers.DiffusionPipeline.load_config(path))
@@ -33,8 +42,7 @@ def load_init_dict(cls: Type[diffusers.DiffusionPipeline], path: os.PathLike):
     R: Dict[str, Tuple[str]] = {}
     for k, v in merged:
         if isinstance(v, list):
-            if v[0] is None or v[1] is None:
-                log.debug(f"Skipping {k} while loading init dict of '{path}': {v}")
+            if k not in cls.__init__.__annotations__:
                 continue
             R[k] = v
     return R
@@ -99,15 +107,8 @@ def patch_kwargs(cls: Type[diffusers.DiffusionPipeline], kwargs: Dict) -> Dict:
     return kwargs
 
 
-def load_pipeline(cls: Type[diffusers.DiffusionPipeline], path: os.PathLike, **kwargs_ort):
+def load_pipeline(cls: Type[diffusers.DiffusionPipeline], path: os.PathLike, **kwargs_ort) -> diffusers.DiffusionPipeline:
     if os.path.isdir(path):
         return cls(**patch_kwargs(cls, load_submodels(path, check_pipeline_sdxl(cls), load_init_dict(cls, path), **kwargs_ort)))
     else:
         return cls.from_single_file(path)
-
-
-def construct_refiner_pipeline(**kwargs):
-    kwargs["text_encoder"] = kwargs["text_encoder_2"]
-    kwargs["tokenizer"] = kwargs["tokenizer_2"]
-    del kwargs["text_encoder_2"], kwargs["tokenizer_2"]
-    return diffusers.OnnxStableDiffusionXLImg2ImgPipeline(**patch_kwargs(diffusers.OnnxStableDiffusionXLImg2ImgPipeline, kwargs))

modules/processing_diffusers.py

@@ -67,7 +67,9 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
             for j in range(len(decoded)):
                 images.save_image(decoded[j], path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix=suffix)
 
-    def diffusers_callback_legacy(step: int, timestep: int, latents: torch.FloatTensor):
+    def diffusers_callback_legacy(step: int, timestep: int, latents: typing.Union[torch.FloatTensor, np.ndarray]):
+        if isinstance(latents, np.ndarray): # latents from Onnx pipelines is ndarray.
+            latents = torch.from_numpy(latents)
         shared.state.sampling_step = step
         shared.state.current_latent = latents
         latents = processing_correction.correction_callback(p, timestep, {'latents': latents})

modules/shared.py

@@ -443,7 +443,8 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), {
 
     "onnx_sep": OptionInfo("<h2>ONNX Runtime</h2>", "", gr.HTML),
     "onnx_execution_provider": OptionInfo(get_default_execution_provider().value, 'Execution Provider', gr.Dropdown, lambda: {"choices": available_execution_providers }),
-    "onnx_cache_converted": OptionInfo(True, 'Cache converted models'),
+    "onnx_cache_converted": OptionInfo(True, 'ONNX cache converted models'),
+    "onnx_unload_base": OptionInfo(False, 'ONNX unload base model when processing refiner'),
 }))
 
 options_templates.update(options_section(('system-paths', "System Paths"), {