diff --git a/modules/olive.py b/modules/olive.py index 1faf7742e..95c52c100 100644 --- a/modules/olive.py +++ b/modules/olive.py @@ -38,12 +38,14 @@ class ENVStore: class OliveOptimizerConfig(ENVStore): from_huggingface_cache: bool + use_fp16_fixed_vae: bool is_sdxl: bool width: int height: int batch_size: int + cross_attention_dim: int time_ids_size: int @@ -218,6 +220,8 @@ def vae_encoder_inputs(_, torch_dtype): def vae_encoder_load(model_name): + if config.use_fp16_fixed_vae: + model_name = "madebyollin/sdxl-vae-fp16-fix" model = diffusers.AutoencoderKL.from_pretrained(model_name, subfolder="vae_encoder" if os.path.isdir(os.path.join(model_name, "vae_encoder")) else "vae", **get_loader_arguments()) model.forward = lambda sample, return_dict: model.encode(sample, return_dict)[0].sample() return model @@ -244,6 +248,8 @@ def vae_decoder_inputs(_, torch_dtype): def vae_decoder_load(model_name): + if config.use_fp16_fixed_vae: + model_name = "madebyollin/sdxl-vae-fp16-fix" model = diffusers.AutoencoderKL.from_pretrained(model_name, subfolder="vae_decoder" if os.path.isdir(os.path.join(model_name, "vae_decoder")) else "vae", **get_loader_arguments()) model.forward = model.decode return model diff --git a/modules/onnx.py b/modules/onnx.py index d97c7f2fc..abc2b00ab 100644 --- a/modules/onnx.py +++ b/modules/onnx.py @@ -5,6 +5,7 @@ import diffusers initialized = False submodels_sd = ("text_encoder", "unet", "vae_encoder", "vae_decoder",) submodels_sdxl = ("text_encoder", "text_encoder_2", "unet", "vae_encoder", "vae_decoder",) +submodels_sdxl_refiner = ("text_encoder_2", "unet", "vae_encoder", "vae_decoder",) class OnnxFakeModule: diff --git a/modules/onnx_pipelines.py b/modules/onnx_pipelines.py index 8188e40ad..0feda0e6d 100644 --- a/modules/onnx_pipelines.py +++ b/modules/onnx_pipelines.py @@ -19,9 +19,10 @@ from installer import log from modules import shared from modules.paths import sd_configs_path 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 -from modules.onnx_utils import check_pipeline_sdxl, load_init_dict, load_submodel, load_submodels, load_pipeline, get_sess_options, patch_kwargs +from modules.onnx import OnnxFakeModule, submodels_sd, submodels_sdxl, submodels_sdxl_refiner +from modules.onnx_utils import check_pipeline_sdxl, 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 @@ -34,6 +35,11 @@ class OnnxPipelineBase(OnnxFakeModule, diffusers.DiffusionPipeline, metaclass=AB def __init__(self): self.model_type = self.__class__.__name__ + def override_processing(self, p: StableDiffusionProcessing): + disable_classifier_free_guidance = p.cfg_scale < 0.01 or "turbo" in self.sd_checkpoint_info.model_name.lower() + if disable_classifier_free_guidance: + p.cfg_scale = 0.0 + @classmethod def from_pretrained(cls, pretrained_model_name_or_path, **_): return OnnxRawPipeline( @@ -56,6 +62,7 @@ class OnnxPipelineBase(OnnxFakeModule, diffusers.DiffusionPipeline, metaclass=AB class OnnxRawPipeline(OnnxPipelineBase): config = {} _is_sdxl: bool + is_refiner: bool from_huggingface_cache: bool path: os.PathLike original_filename: str @@ -69,12 +76,13 @@ class OnnxRawPipeline(OnnxPipelineBase): def __init__(self, constructor: Type[OnnxPipelineBase], path: os.PathLike): self.model_type = constructor.__name__ self._is_sdxl = check_pipeline_sdxl(constructor) + self.is_refiner = "refiner" in str(path).lower() self.from_huggingface_cache = shared.opts.diffusers_dir in os.path.abspath(path) self.path = path self.original_filename = os.path.basename(path) - self.constructor = constructor - self.submodels = submodels_sdxl if self._is_sdxl else submodels_sd + self.constructor = construct_refiner_pipeline if self.is_refiner else constructor + self.submodels = (submodels_sdxl_refiner if self.is_refiner else submodels_sdxl) if self._is_sdxl else submodels_sd if os.path.isdir(path): self.init_dict = load_init_dict(constructor, path) self.scheduler = load_submodel(self.path, None, "scheduler", self.init_dict["scheduler"]) @@ -242,8 +250,11 @@ class OnnxRawPipeline(OnnxPipelineBase): olive_config["input_model"]["config"]["model_path"] = os.path.abspath(os.path.join(in_dir, submodel, "model.onnx")) olive_config["passes"][pass_key]["config"]["float16"] = shared.opts.onnx_olive_float16 olive_config["engine"]["execution_providers"] = [shared.opts.onnx_execution_provider] - if (shared.opts.onnx_execution_provider == ExecutionProvider.CUDA or shared.opts.onnx_execution_provider == ExecutionProvider.ROCm) and version.parse(ort.__version__) < version.parse("1.17.0"): - olive_config["passes"][pass_key]["config"]["optimization_options"] = {"enable_skip_group_norm": False} + if shared.opts.onnx_execution_provider == ExecutionProvider.CUDA or shared.opts.onnx_execution_provider == ExecutionProvider.ROCm: + if version.parse(ort.__version__) < version.parse("1.17.0"): + olive_config["passes"][pass_key]["config"]["optimization_options"] = {"enable_skip_group_norm": False} + if shared.opts.onnx_olive_float16: + olive_config["passes"][pass_key]["config"]["keep_io_types"] = False run(olive_config) @@ -317,28 +328,27 @@ class OnnxRawPipeline(OnnxPipelineBase): shutil.rmtree(out_dir, ignore_errors=True) return None - def preprocess(self, batch_size: int, height: int, width: int): + def preprocess(self, p: StableDiffusionProcessing): + in_dir = self.path if os.path.isdir(self.path) else shared.opts.onnx_temp_dir + config.from_huggingface_cache = self.from_huggingface_cache + config.use_fp16_fixed_vae = self._is_sdxl and not shared.opts.diffusers_vae_upcast config.is_sdxl = self._is_sdxl - config.width = width - config.height = height - config.batch_size = batch_size + config.width = p.width + config.height = p.height + config.batch_size = p.batch_size - if self._is_sdxl: - config.cross_attention_dim = 2048 - config.time_ids_size = 6 - else: - config.cross_attention_dim = height + 256 - config.time_ids_size = 5 + 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 kwargs = { "provider": get_provider(), - "sess_options": get_sess_options(batch_size, height, width, self._is_sdxl), + "sess_options": get_sess_options(p.batch_size if p.cfg_scale < 0.01 or "turbo" in str(self.path).lower() else p.batch_size * 2, p.height, p.width, self._is_sdxl), } - converted_dir = self.convert(self.path if os.path.isdir(self.path) else shared.opts.onnx_temp_dir) + 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)) @@ -346,7 +356,7 @@ class OnnxRawPipeline(OnnxPipelineBase): if shared.opts.onnx_enable_olive: log.warning("Olive implementation is experimental. It contains potentially an issue and is subject to change at any time.") - if width != height: + if p.width != p.height: log.warning("Olive detected different width and height. The quality of the result is not guaranteed.") optimized_dir = self.optimize(converted_dir) if optimized_dir is None: diff --git a/modules/onnx_utils.py b/modules/onnx_utils.py index 16e4181d0..ad2210e67 100644 --- a/modules/onnx_utils.py +++ b/modules/onnx_utils.py @@ -6,20 +6,20 @@ import onnxruntime as ort from installer import log -def get_sess_options(batch_size: int, height: int, width: int, is_sdxl: bool) -> ort.SessionOptions: +def get_sess_options(hidden_batch_size: int, height: int, width: int, is_sdxl: bool) -> ort.SessionOptions: sess_options = ort.SessionOptions() sess_options.enable_mem_pattern = False - sess_options.add_free_dimension_override_by_name("unet_sample_batch", batch_size * 2) + sess_options.add_free_dimension_override_by_name("unet_sample_batch", hidden_batch_size) sess_options.add_free_dimension_override_by_name("unet_sample_channels", 4) sess_options.add_free_dimension_override_by_name("unet_sample_height", height // 8) sess_options.add_free_dimension_override_by_name("unet_sample_width", width // 8) sess_options.add_free_dimension_override_by_name("unet_time_batch", 1) - sess_options.add_free_dimension_override_by_name("unet_hidden_batch", batch_size * 2) + sess_options.add_free_dimension_override_by_name("unet_hidden_batch", hidden_batch_size) sess_options.add_free_dimension_override_by_name("unet_hidden_sequence", 77) if is_sdxl: - sess_options.add_free_dimension_override_by_name("unet_text_embeds_batch", batch_size * 2) + sess_options.add_free_dimension_override_by_name("unet_text_embeds_batch", hidden_batch_size) sess_options.add_free_dimension_override_by_name("unet_text_embeds_size", 1280) - sess_options.add_free_dimension_override_by_name("unet_time_ids_batch", batch_size * 2) + sess_options.add_free_dimension_override_by_name("unet_time_ids_batch", hidden_batch_size) sess_options.add_free_dimension_override_by_name("unet_time_ids_size", 6) return sess_options @@ -76,11 +76,10 @@ def load_submodels(path: os.PathLike, is_sdxl: bool, init_dict: Dict[str, Type], def patch_kwargs(cls: Type[diffusers.DiffusionPipeline], kwargs: Dict) -> Dict: - from modules import onnx_pipelines as pipelines - if cls == pipelines.OnnxStableDiffusionPipeline or cls == pipelines.OnnxStableDiffusionImg2ImgPipeline or cls == pipelines.OnnxStableDiffusionInpaintPipeline: + if cls == diffusers.OnnxStableDiffusionPipeline or cls == diffusers.OnnxStableDiffusionImg2ImgPipeline or cls == diffusers.OnnxStableDiffusionInpaintPipeline: kwargs["safety_checker"] = None kwargs["requires_safety_checker"] = False - if cls == pipelines.OnnxStableDiffusionXLPipeline or cls == pipelines.OnnxStableDiffusionXLImg2ImgPipeline: + if cls == diffusers.OnnxStableDiffusionXLPipeline or cls == diffusers.OnnxStableDiffusionXLImg2ImgPipeline: kwargs["config"] = {} return kwargs @@ -91,3 +90,10 @@ def load_pipeline(cls: Type[diffusers.DiffusionPipeline], path: os.PathLike, **k 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)) diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 999274e83..287ea6e15 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -22,7 +22,10 @@ def process_diffusers(p: processing.StableDiffusionProcessing): results = [] if hasattr(shared.sd_model, 'preprocess'): - shared.sd_model = shared.sd_model.preprocess(p.batch_size, p.height, p.width) + shared.sd_model = shared.sd_model.preprocess(p) + + if hasattr(shared.sd_model, 'override_processing'): + shared.sd_model.override_processing(p) def is_txt2img(): return sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE