From def8a300173865bb15eb7b0a15b27942b55cea8b Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Sat, 21 Oct 2023 17:47:39 +0900 Subject: [PATCH] Initial Olive implementation. --- configs/olive/config_text_encoder.json | 93 +++++ configs/olive/config_unet.json | 113 ++++++ configs/olive/config_vae_decoder.json | 100 +++++ configs/olive/config_vae_encoder.json | 100 +++++ installer.py | 2 + modules/olive.py | 488 +++++++++++++++++++++++++ modules/paths.py | 124 +++---- modules/processing_diffusers.py | 6 +- modules/shared.py | 5 + modules/shared_items.py | 2 + requirements.txt | 35 +- 11 files changed, 977 insertions(+), 91 deletions(-) create mode 100644 configs/olive/config_text_encoder.json create mode 100644 configs/olive/config_unet.json create mode 100644 configs/olive/config_vae_decoder.json create mode 100644 configs/olive/config_vae_encoder.json create mode 100644 modules/olive.py diff --git a/configs/olive/config_text_encoder.json b/configs/olive/config_text_encoder.json new file mode 100644 index 000000000..fb3d8feb0 --- /dev/null +++ b/configs/olive/config_text_encoder.json @@ -0,0 +1,93 @@ +{ + "input_model": { + "type": "PyTorchModel", + "config": { + "model_path": "", + "model_loader": "text_encoder_load", + "model_script": "modules/olive.py", + "io_config": { + "input_names": ["input_ids"], + "output_names": ["last_hidden_state", "pooler_output"], + "dynamic_axes": { "input_ids": { "0": "batch", "1": "sequence" } } + }, + "dummy_inputs_func": "text_encoder_conversion_inputs" + } + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "config": { + "accelerators": ["gpu"] + } + } + }, + "evaluators": { + "common_evaluator": { + "metrics": [ + { + "name": "latency", + "type": "latency", + "sub_types": [{ "name": "avg" }], + "user_config": { + "user_script": "modules/olive.py", + "dataloader_func": "text_encoder_data_loader", + "batch_size": 1 + } + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "config": { + "target_opset": 14 + } + }, + "optimize": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "clip", + "float16": true, + "use_gpu": true, + "keep_io_types": false, + "optimization_options": { + "enable_gelu": true, + "enable_layer_norm": true, + "enable_attention": true, + "use_multi_head_attention": true, + "enable_skip_layer_norm": false, + "enable_embed_layer_norm": true, + "enable_bias_skip_layer_norm": false, + "enable_bias_gelu": true, + "enable_gelu_approximation": false, + "enable_qordered_matmul": false, + "enable_shape_inference": true, + "enable_gemm_fast_gelu": false, + "enable_nhwc_conv": false, + "enable_group_norm": true, + "enable_bias_splitgelu": false, + "enable_packed_qkv": true, + "enable_packed_kv": true, + "enable_bias_add": false + }, + "force_fp32_ops": ["RandomNormalLike"] + } + } + }, + "engine": { + "search_strategy": { + "execution_order": "joint", + "search_algorithm": "exhaustive" + }, + "evaluator": "common_evaluator", + "evaluate_input_model": false, + "host": "local_system", + "target": "local_system", + "cache_dir": "cache", + "output_name": "text_encoder", + "output_dir": "footprints", + "execution_providers": ["DmlExecutionProvider"] + } +} diff --git a/configs/olive/config_unet.json b/configs/olive/config_unet.json new file mode 100644 index 000000000..7e00506d6 --- /dev/null +++ b/configs/olive/config_unet.json @@ -0,0 +1,113 @@ +{ + "input_model": { + "type": "PyTorchModel", + "config": { + "model_path": "", + "model_loader": "unet_load", + "model_script": "modules/olive.py", + "io_config": { + "input_names": [ + "sample", + "timestep", + "encoder_hidden_states", + "return_dict" + ], + "output_names": ["out_sample"], + "dynamic_axes": { + "sample": { + "0": "unet_sample_batch", + "1": "unet_sample_channels", + "2": "unet_sample_height", + "3": "unet_sample_width" + }, + "timestep": { "0": "unet_time_batch" }, + "encoder_hidden_states": { + "0": "unet_hidden_batch", + "1": "unet_hidden_sequence" + } + } + }, + "dummy_inputs_func": "unet_conversion_inputs" + } + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "config": { + "accelerators": ["gpu"] + } + } + }, + "evaluators": { + "common_evaluator": { + "metrics": [ + { + "name": "latency", + "type": "latency", + "sub_types": [{ "name": "avg" }], + "user_config": { + "user_script": "modules/olive.py", + "dataloader_func": "unet_data_loader", + "batch_size": 2 + } + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "config": { + "target_opset": 14, + "save_as_external_data": true, + "all_tensors_to_one_file": true, + "external_data_name": "weights.pb" + } + }, + "optimize": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "unet", + "float16": true, + "use_gpu": true, + "keep_io_types": false, + "optimization_options": { + "enable_gelu": true, + "enable_layer_norm": true, + "enable_attention": true, + "use_multi_head_attention": true, + "enable_skip_layer_norm": false, + "enable_embed_layer_norm": true, + "enable_bias_skip_layer_norm": false, + "enable_bias_gelu": true, + "enable_gelu_approximation": false, + "enable_qordered_matmul": false, + "enable_shape_inference": true, + "enable_gemm_fast_gelu": false, + "enable_nhwc_conv": false, + "enable_group_norm": true, + "enable_bias_splitgelu": false, + "enable_packed_qkv": true, + "enable_packed_kv": true, + "enable_bias_add": false + }, + "force_fp32_ops": ["RandomNormalLike"] + } + } + }, + "engine": { + "search_strategy": { + "execution_order": "joint", + "search_algorithm": "exhaustive" + }, + "evaluator": "common_evaluator", + "evaluate_input_model": false, + "host": "local_system", + "target": "local_system", + "cache_dir": "cache", + "output_name": "unet", + "output_dir": "footprints", + "execution_providers": ["DmlExecutionProvider"] + } +} diff --git a/configs/olive/config_vae_decoder.json b/configs/olive/config_vae_decoder.json new file mode 100644 index 000000000..d7d746d61 --- /dev/null +++ b/configs/olive/config_vae_decoder.json @@ -0,0 +1,100 @@ +{ + "input_model": { + "type": "PyTorchModel", + "config": { + "model_path": "", + "model_loader": "vae_decoder_load", + "model_script": "modules/olive.py", + "io_config": { + "input_names": ["latent_sample", "return_dict"], + "output_names": ["sample"], + "dynamic_axes": { + "latent_sample": { + "0": "batch", + "1": "channels", + "2": "height", + "3": "width" + } + } + }, + "dummy_inputs_func": "vae_decoder_conversion_inputs" + } + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "config": { + "accelerators": ["gpu"] + } + } + }, + "evaluators": { + "common_evaluator": { + "metrics": [ + { + "name": "latency", + "type": "latency", + "sub_types": [{ "name": "avg" }], + "user_config": { + "user_script": "modules/olive.py", + "dataloader_func": "vae_decoder_data_loader", + "batch_size": 1 + } + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "config": { + "target_opset": 14 + } + }, + "optimize": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "vae", + "float16": true, + "use_gpu": true, + "keep_io_types": false, + "optimization_options": { + "enable_gelu": true, + "enable_layer_norm": true, + "enable_attention": true, + "use_multi_head_attention": true, + "enable_skip_layer_norm": false, + "enable_embed_layer_norm": true, + "enable_bias_skip_layer_norm": false, + "enable_bias_gelu": true, + "enable_gelu_approximation": false, + "enable_qordered_matmul": false, + "enable_shape_inference": true, + "enable_gemm_fast_gelu": false, + "enable_nhwc_conv": false, + "enable_group_norm": true, + "enable_bias_splitgelu": false, + "enable_packed_qkv": true, + "enable_packed_kv": true, + "enable_bias_add": false + }, + "force_fp32_ops": ["RandomNormalLike"] + } + } + }, + "engine": { + "search_strategy": { + "execution_order": "joint", + "search_algorithm": "exhaustive" + }, + "evaluator": "common_evaluator", + "evaluate_input_model": false, + "host": "local_system", + "target": "local_system", + "cache_dir": "cache", + "output_name": "vae_decoder", + "output_dir": "footprints", + "execution_providers": ["DmlExecutionProvider"] + } +} diff --git a/configs/olive/config_vae_encoder.json b/configs/olive/config_vae_encoder.json new file mode 100644 index 000000000..f28fe3574 --- /dev/null +++ b/configs/olive/config_vae_encoder.json @@ -0,0 +1,100 @@ +{ + "input_model": { + "type": "PyTorchModel", + "config": { + "model_path": "", + "model_loader": "vae_encoder_load", + "model_script": "modules/olive.py", + "io_config": { + "input_names": ["sample", "return_dict"], + "output_names": ["latent_sample"], + "dynamic_axes": { + "sample": { + "0": "batch", + "1": "channels", + "2": "height", + "3": "width" + } + } + }, + "dummy_inputs_func": "vae_encoder_conversion_inputs" + } + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "config": { + "accelerators": ["gpu"] + } + } + }, + "evaluators": { + "common_evaluator": { + "metrics": [ + { + "name": "latency", + "type": "latency", + "sub_types": [{ "name": "avg" }], + "user_config": { + "user_script": "modules/olive.py", + "dataloader_func": "vae_encoder_data_loader", + "batch_size": 1 + } + } + ] + } + }, + "passes": { + "convert": { + "type": "OnnxConversion", + "config": { + "target_opset": 14 + } + }, + "optimize": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "vae", + "float16": true, + "use_gpu": true, + "keep_io_types": false, + "optimization_options": { + "enable_gelu": true, + "enable_layer_norm": true, + "enable_attention": true, + "use_multi_head_attention": true, + "enable_skip_layer_norm": false, + "enable_embed_layer_norm": true, + "enable_bias_skip_layer_norm": false, + "enable_bias_gelu": true, + "enable_gelu_approximation": false, + "enable_qordered_matmul": false, + "enable_shape_inference": true, + "enable_gemm_fast_gelu": false, + "enable_nhwc_conv": false, + "enable_group_norm": true, + "enable_bias_splitgelu": false, + "enable_packed_qkv": true, + "enable_packed_kv": true, + "enable_bias_add": false + }, + "force_fp32_ops": ["RandomNormalLike"] + } + } + }, + "engine": { + "search_strategy": { + "execution_order": "joint", + "search_algorithm": "exhaustive" + }, + "evaluator": "common_evaluator", + "evaluate_input_model": false, + "host": "local_system", + "target": "local_system", + "cache_dir": "cache", + "output_name": "vae_encoder", + "output_dir": "footprints", + "execution_providers": ["DmlExecutionProvider"] + } +} diff --git a/installer.py b/installer.py index 14a16847d..d5517b915 100644 --- a/installer.py +++ b/installer.py @@ -583,6 +583,8 @@ def install_packages(): install(clip_package, 'clip') invisiblewatermark_package = os.environ.get('INVISIBLEWATERMARK_PACKAGE', "git+https://github.com/patrickvonplaten/invisible-watermark.git@remove_onnxruntime_depedency") install(invisiblewatermark_package, 'invisible-watermark') + install('olive-ai[directml]', 'olive-ai', ignore=True) + install('onnxruntime-directml==1.16.1', 'onnxruntime-directml', ignore=True) install('pi-heif', 'pi_heif', ignore=True) tensorflow_package = os.environ.get('TENSORFLOW_PACKAGE', 'tensorflow==2.13.0') install(tensorflow_package, 'tensorflow-rocm' if 'rocm' in tensorflow_package else 'tensorflow', ignore=True) diff --git a/modules/olive.py b/modules/olive.py new file mode 100644 index 000000000..88a80ebf7 --- /dev/null +++ b/modules/olive.py @@ -0,0 +1,488 @@ +import os +import json +import torch +import shutil +import diffusers +import numpy as np +from typing import Union, Optional, Callable, List +from transformers.models.clip.modeling_clip import CLIPTextModel, CLIPTextModelWithProjection +from installer import log, args +from modules.shared import opts, cmd_opts +from modules.paths import models_path, sd_configs_path +from modules.sd_models import CheckpointInfo + +temp_dir = os.path.join(models_path, "OliveTemp") +cache_dir = os.path.join(models_path, "OliveCache") + +submodels = ("text_encoder", "unet", "vae_encoder", "vae_decoder",) + +execution_provider = "CUDAExecutionProvider" +if args.use_directml: + execution_provider = "DmlExecutionProvider" +elif args.use_rocm: + execution_provider = "ROCmExecutionProvider" +provider = (execution_provider, { + "device_id": int(cmd_opts.device_id or 0), +}) + +class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline): + sd_model_hash: str + sd_checkpoint_info: CheckpointInfo + sd_model_checkpoint: str + + def apply(self, dummy_pipeline): + self.sd_model_hash = dummy_pipeline.sd_model_hash + self.sd_checkpoint_info = dummy_pipeline.sd_checkpoint_info + self.sd_model_checkpoint = dummy_pipeline.sd_model_checkpoint + return self + + def __call__( + self, + prompt: Union[str, List[str]] = None, + height: Optional[int] = 512, + width: Optional[int] = 512, + num_inference_steps: Optional[int] = 50, + guidance_scale: Optional[float] = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: Optional[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, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback: Optional[Callable[[int, int, np.ndarray], None]] = None, + callback_steps: int = 1, + ): + import inspect + + # check inputs. Raise error if not correct + self.check_inputs( + prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds + ) + + # define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + # 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 + + 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, + ) + + # get the initial random noise unless the user supplied it + latents_dtype = prompt_embeds.dtype + latents_shape = (batch_size * num_images_per_prompt, 4, height // 8, width // 8) + 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(latents_shape).astype(latents_dtype) + elif latents.shape != latents_shape: + raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") + + # set timesteps + self.scheduler.set_timesteps(num_inference_steps) + + latents = latents * np.float64(self.scheduler.init_noise_sigma) + + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs["eta"] = eta + + timestep_dtype = next( + (input.type for input in self.unet.model.get_inputs() if input.name == "timestep"), "tensor(float)" + ) + timestep_dtype = diffusers.pipelines.onnx_utils.ORT_TO_NP_TYPE[timestep_dtype] + + for i, t in enumerate(self.progress_bar(self.scheduler.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) + 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) + + # 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 callback is not None and i % callback_steps == 0: + callback(i, t, torch.from_numpy(latents)) + + latents = 1 / 0.18215 * latents + + has_nsfw_concept = None + + if not output_type == "latent": + # image = self.vae_decoder(latent_sample=latents)[0] + # 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])] + ) + + image = np.clip(image / 2 + 0.5, 0, 1) + image = image.transpose((0, 2, 3, 1)) + + if self.safety_checker is not None: + safety_checker_input = self.feature_extractor( + self.numpy_to_pil(image), return_tensors="np" + ).pixel_values.astype(image.dtype) + + images, has_nsfw_concept = [], [] + for i in range(image.shape[0]): + image_i, has_nsfw_concept_i = self.safety_checker( + clip_input=safety_checker_input[i : i + 1], images=image[i : i + 1] + ) + images.append(image_i) + has_nsfw_concept.append(has_nsfw_concept_i[0]) + image = np.concatenate(images) + + if output_type == "pil": + image = self.numpy_to_pil(image) + else: + image = latents + + if not return_dict: + return (image, has_nsfw_concept) + + return diffusers.pipelines.stable_diffusion.StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) + +class OlivePipeline(diffusers.DiffusionPipeline): + sd_model_hash: str + sd_checkpoint_info: CheckpointInfo + sd_model_checkpoint: str + + unoptimized: diffusers.DiffusionPipeline + original_filename: str + + def __init__(self, path, pipeline: diffusers.DiffusionPipeline, scheduler: diffusers.schedulers.scheduling_utils.KarrasDiffusionSchedulers = 0): + self.original_filename = os.path.basename(path) + self.unoptimized = pipeline + del pipeline + if not os.path.exists(temp_dir): + os.mkdir(temp_dir) + self.unoptimized.save_pretrained(temp_dir) + + @staticmethod + def from_pretrained(pretrained_model_name_or_path, **kwargs): + return OlivePipeline(pretrained_model_name_or_path, diffusers.DiffusionPipeline.from_pretrained(pretrained_model_name_or_path, **kwargs)) + + @staticmethod + def from_single_file(pretrained_model_name_or_path, **kwargs): + return OlivePipeline(pretrained_model_name_or_path, diffusers.StableDiffusionPipeline.from_single_file(pretrained_model_name_or_path, **kwargs)) + + @staticmethod + def from_ckpt(*args, **kwargs): + return OlivePipeline.from_single_file(**args, **kwargs) + + def to(self, *args, **kwargs): + pass + + def optimize(self, width: int, height: int): + from olive.workflows import run + from olive.model import ONNXModel + + out_dir = os.path.join(cache_dir, f"{self.original_filename}-{width}w-{height}h") + if os.path.isdir(out_dir): + del self.unoptimized + return OnnxStableDiffusionPipeline.from_pretrained(out_dir, provider=provider).apply(self) + + try: + shutil.copytree( + temp_dir, out_dir, ignore=shutil.ignore_patterns("weights.pb", "*.onnx", "*.safetensors", "*.ckpt") + ) + + optimize_config["width"] = width + optimize_config["height"] = height + + optimized_model_paths = {} + + for submodel in submodels: + log.info(f"\nOptimizing {submodel}") + + with open(os.path.join(sd_configs_path, "olive", f"config_{submodel}.json"), "r") as config_file: + olive_config = json.load(config_file) + olive_config["passes"]["optimize"]["config"]["float16"] = opts.olive_float16 + + run(olive_config) + + with open(os.path.join("footprints", f"{submodel}_gpu-dml_footprints.json"), "r") as footprint_file: + footprints = json.load(footprint_file) + conversion_footprint = None + optimizer_footprint = None + for _, footprint in footprints.items(): + if footprint["from_pass"] == "OnnxConversion": + conversion_footprint = footprint + elif footprint["from_pass"] == "OrtTransformersOptimization": + optimizer_footprint = footprint + + assert conversion_footprint and optimizer_footprint, "Failed to optimize model" + + optimized_model_paths[submodel] = ONNXModel( + **optimizer_footprint["model_config"]["config"] + ).model_path + + log.info(f"Optimized {submodel}") + shutil.rmtree("footprints") + shutil.rmtree(temp_dir) + + kwargs = { + "tokenizer": self.unoptimized.tokenizer, + "scheduler": self.unoptimized.scheduler, + "safety_checker": self.unoptimized.safety_checker if hasattr(self.unoptimized, "safety_checker") else None, + "feature_extractor": self.unoptimized.feature_extractor, + } + del self.unoptimized + for submodel in submodels: + kwargs[submodel] = diffusers.OnnxRuntimeModel.from_pretrained( + os.path.dirname(optimized_model_paths[submodel]), provider=provider, + ) + + pipeline = OnnxStableDiffusionPipeline( + **kwargs, + requires_safety_checker=False, + ).apply(self) + pipeline.to_json_file(os.path.join(out_dir, "model_index.json")) + del kwargs + + for submodel in submodels: + src_path = optimized_model_paths[submodel] + src_parent = os.path.dirname(src_path) + dst_parent = os.path.join(out_dir, submodel) + dst_path = os.path.join(dst_parent, "model.onnx") + if not os.path.isdir(dst_parent): + os.mkdir(dst_parent) + shutil.copyfile(src_path, dst_path) + + weights_src_path = os.path.join(src_parent, (os.path.basename(src_path) + ".data")) + if os.path.isfile(weights_src_path): + weights_dst_path = os.path.join(dst_parent, (os.path.basename(dst_path) + ".data")) + shutil.copyfile(weights_src_path, weights_dst_path) + return pipeline + except Exception as e: + log.error(f"Failed to optimize model '{self.original_filename}'.") + log.error(e) + shutil.rmtree("cache", ignore_errors=True) + shutil.rmtree("footprints", ignore_errors=True) + shutil.rmtree(temp_dir, ignore_errors=True) + shutil.rmtree(out_dir, ignore_errors=True) + return self.unoptimized + +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +optimize_config = { + "is_sdxl": False, + + "source": os.path.abspath(temp_dir), + + "width": 512, + "height": 512, +} + + +# Helper latency-only dataloader that creates random tensors with no label +class RandomDataLoader: + def __init__(self, create_inputs_func, batchsize, torch_dtype): + self.create_input_func = create_inputs_func + self.batchsize = batchsize + self.torch_dtype = torch_dtype + + def __getitem__(self, idx): + label = None + return self.create_input_func(self.batchsize, self.torch_dtype), label + +# ----------------------------------------------------------------------------- +# TEXT ENCODER +# ----------------------------------------------------------------------------- + + +def text_encoder_inputs(batchsize, torch_dtype): + input_ids = torch.zeros((batchsize, 77), dtype=torch_dtype) + return { + "input_ids": input_ids, + "output_hidden_states": True, + } if optimize_config["is_sdxl"] else input_ids + + +def text_encoder_load(model_name): + model = CLIPTextModel.from_pretrained(optimize_config["source"], subfolder="text_encoder") + return model + + +def text_encoder_conversion_inputs(model): + return text_encoder_inputs(1, torch.int32) + + +def text_encoder_data_loader(data_dir, batchsize, *args, **kwargs): + return RandomDataLoader(text_encoder_inputs, batchsize, torch.int32) + + +# ----------------------------------------------------------------------------- +# TEXT ENCODER 2 +# ----------------------------------------------------------------------------- + + +def text_encoder_2_inputs(batchsize, torch_dtype): + return { + "input_ids": torch.zeros((batchsize, 77), dtype=torch_dtype), + "output_hidden_states": True, + } + + +def text_encoder_2_load(model_name): + model = CLIPTextModelWithProjection.from_pretrained(optimize_config["source"], subfolder="text_encoder_2") + return model + + +def text_encoder_2_conversion_inputs(model): + return text_encoder_2_inputs(1, torch.int64) + + +def text_encoder_2_data_loader(data_dir, batchsize, *args, **kwargs): + return RandomDataLoader(text_encoder_2_inputs, batchsize, torch.int64) + + +# ----------------------------------------------------------------------------- +# UNET +# ----------------------------------------------------------------------------- + + +def unet_inputs(batchsize, torch_dtype, is_conversion_inputs=False): + # TODO (pavignol): All the multiplications by 2 here are bacause the XL base has 2 text encoders + # For refiner, it should be multiplied by 1 (single text encoder) + height = optimize_config["height"] + width = optimize_config["width"] + + if optimize_config["is_sdxl"]: + inputs = { + "sample": torch.rand((2 * batchsize, 4, height // 8, width // 8), dtype=torch_dtype), + "timestep": torch.rand((1,), dtype=torch_dtype), + "encoder_hidden_states": torch.rand((2 * batchsize, 77, height * 2), dtype=torch_dtype), + } + + if is_conversion_inputs: + inputs["additional_inputs"] = { + "added_cond_kwargs": { + "text_embeds": torch.rand((2 * batchsize, height + 256), dtype=torch_dtype), + "time_ids": torch.rand((2 * batchsize, 6), dtype=torch_dtype), + } + } + else: + inputs["text_embeds"] = torch.rand((2 * batchsize, height + 256), dtype=torch_dtype) + inputs["time_ids"] = torch.rand((2 * batchsize, 6), dtype=torch_dtype) + else: + inputs = { + "sample": torch.rand((batchsize, 4, height // 8, width // 8), dtype=torch_dtype), + "timestep": torch.rand((batchsize,), dtype=torch_dtype), + "encoder_hidden_states": torch.rand((batchsize, 77, height + 256), dtype=torch_dtype), + "return_dict": False, + } + + return inputs + + +def unet_load(model_name): + model = diffusers.UNet2DConditionModel.from_pretrained(optimize_config["source"], subfolder="unet") + return model + + +def unet_conversion_inputs(model): + return tuple(unet_inputs(1, torch.float32, True).values()) + + +def unet_data_loader(data_dir, batchsize, *args, **kwargs): + return RandomDataLoader(unet_inputs, batchsize, torch.float16) + + +# ----------------------------------------------------------------------------- +# VAE ENCODER +# ----------------------------------------------------------------------------- + + +def vae_encoder_inputs(batchsize, torch_dtype): + return { + "sample": torch.rand((batchsize, 3, optimize_config["height"], optimize_config["width"]), dtype=torch_dtype), + "return_dict": False, + } + + +def vae_encoder_load(model_name): + source = os.path.join(optimize_config["source"], "vae") + if not os.path.isdir(source): + source += "_encoder" + model = diffusers.AutoencoderKL.from_pretrained(source) + model.forward = lambda sample, return_dict: model.encode(sample, return_dict)[0].sample() + return model + + +def vae_encoder_conversion_inputs(model): + return tuple(vae_encoder_inputs(1, torch.float32).values()) + + +def vae_encoder_data_loader(data_dir, batchsize, *args, **kwargs): + return RandomDataLoader(vae_encoder_inputs, batchsize, torch.float16) + + +# ----------------------------------------------------------------------------- +# VAE DECODER +# ----------------------------------------------------------------------------- + + +def vae_decoder_inputs(batchsize, torch_dtype): + return { + "latent_sample": torch.rand((batchsize, 4, optimize_config["height"] // 8, optimize_config["width"] // 8), dtype=torch_dtype), + "return_dict": False, + } + + +def vae_decoder_load(model_name): + source = os.path.join(optimize_config["source"], "vae") + if not os.path.isdir(source): + source += "_decoder" + model = diffusers.AutoencoderKL.from_pretrained(source) + model.forward = model.decode + return model + + +def vae_decoder_conversion_inputs(model): + return tuple(vae_decoder_inputs(1, torch.float32).values()) + + +def vae_decoder_data_loader(data_dir, batchsize, *args, **kwargs): + return RandomDataLoader(vae_decoder_inputs, batchsize, torch.float16) diff --git a/modules/paths.py b/modules/paths.py index 7295f6575..60c69553b 100644 --- a/modules/paths.py +++ b/modules/paths.py @@ -1,79 +1,65 @@ -# this module must not have any dependencies as it is a very first import before webui even starts import os import sys -import json -import argparse -from modules.errors import log +import olive.workflows +from modules import paths_internal, errors -# parse args, parse again after we have the data-dir and early-read the config file -parser = argparse.ArgumentParser(add_help=False) -parser.add_argument("--ckpt", type=str, default=os.environ.get("SD_MODEL", None), help="Path to model checkpoint to load immediately, default: %(default)s") -parser.add_argument("--data-dir", type=str, default=os.environ.get("SD_DATADIR", ''), help="Base path where all user data is stored, default: %(default)s") -parser.add_argument("--models-dir", type=str, default=os.environ.get("SD_MODELSDIR", None), help="Base path where all models are stored, default: %(default)s",) -cli = parser.parse_known_args()[0] -parser.add_argument("--config", type=str, default=os.environ.get("SD_CONFIG", os.path.join(cli.data_dir, 'config.json')), help="Use specific server configuration file, default: %(default)s") -cli = parser.parse_known_args()[0] -config_path = cli.config if os.path.isabs(cli.config) else os.path.join(cli.data_dir, cli.config) -try: - with open(config_path, 'r', encoding='utf8') as f: - config = json.load(f) -except Exception: - config = {} +debug = errors.log.info if os.environ.get('SD_PATH_DEBUG', None) is not None else lambda *args, **kwargs: None +data_path = paths_internal.data_path +script_path = paths_internal.script_path +models_path = paths_internal.models_path +sd_configs_path = paths_internal.sd_configs_path +sd_default_config = paths_internal.sd_default_config +sd_model_file = paths_internal.sd_model_file +default_sd_model_file = paths_internal.default_sd_model_file +extensions_dir = paths_internal.extensions_dir +extensions_builtin_dir = paths_internal.extensions_builtin_dir + +# data_path = cmd_opts_pre.data +sys.path.insert(0, script_path) + +# search for directory of stable diffusion in following places +sd_path = None +possible_sd_paths = [os.path.join(script_path, 'repositories/stable-diffusion-stability-ai'), '.', os.path.dirname(script_path)] +for possible_sd_path in possible_sd_paths: + if os.path.exists(os.path.join(possible_sd_path, 'ldm/models/diffusion/ddpm.py')): + sd_path = os.path.abspath(possible_sd_path) + break + +assert sd_path is not None, f"Couldn't find Stable Diffusion in any of: {possible_sd_paths}" + +path_dirs = [ + (sd_path, 'ldm', 'Stable Diffusion', []), + (os.path.join(sd_path, '../taming-transformers'), 'taming', 'Taming Transformers', []), + (os.path.join(sd_path, '../CodeFormer'), 'inference_codeformer.py', 'CodeFormer', []), + (os.path.join(sd_path, '../BLIP'), 'models/blip.py', 'BLIP', []), + (os.path.join(sd_path, '../k-diffusion'), 'k_diffusion/sampling.py', 'k_diffusion', ["atstart"]), +] -modules_path = os.path.dirname(os.path.realpath(__file__)) -script_path = os.path.dirname(modules_path) -data_path = cli.data_dir -models_config = cli.models_dir or config.get('models_dir') or 'models' -models_path = models_config if os.path.isabs(models_config) else os.path.join(data_path, models_config) -extensions_dir = os.path.join(data_path, "extensions") -extensions_builtin_dir = "extensions-builtin" -sd_configs_path = os.path.join(script_path, "configs") -sd_default_config = os.path.join(sd_configs_path, "v1-inference.yaml") -sd_model_file = cli.ckpt or os.path.join(script_path, 'model.ckpt') # not used -default_sd_model_file = sd_model_file # not used -debug = log.trace if os.environ.get('SD_PATH_DEBUG', None) is not None else lambda *args, **kwargs: None -debug('Trace: PATH') paths = {} -if os.environ.get('SD_PATH_DEBUG', None) is not None: - print(f'Paths: script-path="{script_path}" data-dir="{data_path}" models-dir="{models_path}" config="{config_path}"') - - -def register_paths(): - log.debug('Register paths') - sys.path.insert(0, script_path) - sd_path = os.path.join(script_path, 'repositories') - path_dirs = [ - (sd_path, 'ldm', 'ldm', []), - (sd_path, 'taming', 'Taming Transformers', []), - (os.path.join(sd_path, 'blip'), 'models/blip.py', 'BLIP', []), - (os.path.join(sd_path, 'codeformer'), 'inference_codeformer.py', 'CodeFormer', []), - (os.path.join(modules_path, 'k-diffusion'), 'k_diffusion/sampling.py', 'k_diffusion', ["atstart"]), - ] - for d, must_exist, what, _options in path_dirs: - must_exist_path = os.path.abspath(os.path.join(script_path, d, must_exist)) - if not os.path.exists(must_exist_path): - log.error(f'Required path not found: path={must_exist_path} item={what}') - else: - d = os.path.abspath(d) - sys.path.append(d) - paths[what] = d - - -def create_path(folder): - if folder is None or folder == '': - return - if os.path.exists(folder): - return - try: - os.makedirs(folder, exist_ok=True) - log.info(f'Create: folder="{folder}"') - except Exception as e: - log.error(f'Create failed: folder="{folder}" {e}') +for d, must_exist, what, _options in path_dirs: + must_exist_path = os.path.abspath(os.path.join(script_path, d, must_exist)) + if not os.path.exists(must_exist_path): + errors.log.error(f'Required path not found: path={must_exist_path} item={what}') + else: + d = os.path.abspath(d) + sys.path.append(d) + paths[what] = d def create_paths(opts): + def create_path(folder): + if folder is None or folder == '': + return + if os.path.exists(folder): + return + try: + os.makedirs(folder, exist_ok=True) + errors.log.info(f'Create folder={folder}') + except Exception as e: + errors.log.error(f'Create Failed folder={folder} {e}') + def fix_path(folder): tgt = opts.data.get(folder, None) or opts.data_labels[folder].default if tgt is None or tgt == '': @@ -85,7 +71,7 @@ def create_paths(opts): fix = os.path.abspath(fix) fix = fix if os.path.isabs(fix) else os.path.relpath(fix, script_path) opts.data[folder] = fix - debug(f'Paths: folder="{folder}" original="{tgt}" target="{fix}"') + debug(f'Paths: folder={folder} original="{tgt}" target="{fix}"') return opts.data[folder] create_path(data_path) @@ -104,15 +90,11 @@ def create_paths(opts): create_path(fix_path('outdir_samples')) create_path(fix_path('outdir_txt2img_samples')) create_path(fix_path('outdir_img2img_samples')) - create_path(fix_path('outdir_control_samples')) create_path(fix_path('outdir_extras_samples')) - create_path(fix_path('outdir_init_images')) create_path(fix_path('outdir_grids')) create_path(fix_path('outdir_txt2img_grids')) create_path(fix_path('outdir_img2img_grids')) - create_path(fix_path('outdir_control_grids')) create_path(fix_path('outdir_save')) - create_path(fix_path('outdir_video')) create_path(fix_path('styles_dir')) diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index b8e58d828..c7c58e0a6 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -7,6 +7,7 @@ import torch import torchvision.transforms.functional as TF import diffusers from modules import shared, devices, processing, sd_samplers, sd_models, images, errors, masking, prompt_parser_diffusers, sd_hijack_hypertile, processing_correction, processing_vae +from modules.olive import OlivePipeline debug = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -20,6 +21,9 @@ def process_diffusers(p: processing.StableDiffusionProcessing): orig_pipeline = shared.sd_model results = [] + if isinstance(shared.sd_model, OlivePipeline): + shared.sd_model = shared.sd_model.optimize(p.width, p.height) + def is_txt2img(): return sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE @@ -219,7 +223,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing): generator = [torch.Generator(generator_device).manual_seed(s) for s in p.seeds] prompts, negative_prompts, prompts_2, negative_prompts_2 = fix_prompts(prompts, negative_prompts, prompts_2, negative_prompts_2) parser = 'Fixed attention' - if shared.opts.prompt_attention != 'Fixed attention' and 'StableDiffusion' in model.__class__.__name__: + if shared.opts.prompt_attention != 'Fixed attention' and 'StableDiffusion' in model.__class__.__name__ and not isinstance(model, diffusers.OnnxStableDiffusionPipeline): try: prompt_parser_diffusers.encode_prompts(model, p, prompts, negative_prompts, kwargs.get("num_inference_steps", 1), kwargs.pop("clip_skip", None)) parser = shared.opts.prompt_attention diff --git a/modules/shared.py b/modules/shared.py index 9dffea095..04a55e385 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -434,6 +434,11 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), { "diffusers_aesthetics_score": OptionInfo(False, "Require aesthetics score"), "diffusers_pooled": OptionInfo("default", "Diffusers SDXL pooled embeds", gr.Radio, {"choices": ['default', 'weighted']}), "huggingface_token": OptionInfo('', 'HuggingFace token'), + + "olive_sep": OptionInfo("

Olive

", "", gr.HTML), + "olive_float16": OptionInfo(True, 'Use FP16 (will use FP32 if unchecked)'), + "olive_cache_optimized": OptionInfo(False, 'Cache optimized models'), + "olive_garbage_collect": OptionInfo(False, 'Collect garbage at the end of each generation'), })) options_templates.update(options_section(('system-paths', "System Paths"), { diff --git a/modules/shared_items.py b/modules/shared_items.py index 41952241a..1b004dad8 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -26,6 +26,7 @@ def list_crossattention(): def get_pipelines(): import diffusers + from modules.olive import OlivePipeline from installer import log pipelines = { # note: not all pipelines can be used manually as they require prior pipeline next to decoder pipeline 'Autodetect': None, @@ -38,6 +39,7 @@ def get_pipelines(): 'Stable Diffusion XL Img2Img': getattr(diffusers, 'StableDiffusionXLImg2ImgPipeline', None), 'Stable Diffusion XL Inpaint': getattr(diffusers, 'StableDiffusionXLInpaintPipeline', None), 'Stable Diffusion XL Instruct': getattr(diffusers, 'StableDiffusionXLInstructPix2PixPipeline', None), + 'ONNX Stable Diffusion (Olive)': OlivePipeline, 'Latent Consistency Model': getattr(diffusers, 'LatentConsistencyModelPipeline', None), 'PixArt Alpha': getattr(diffusers, 'PixArtAlphaPipeline', None), 'UniDiffuser': getattr(diffusers, 'UniDiffuserPipeline', None), diff --git a/requirements.txt b/requirements.txt index db1ffee1a..6af77209b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -24,6 +24,7 @@ lmdb lpips omegaconf open-clip-torch +opencv-contrib-python-headless piexif psutil pyyaml @@ -32,7 +33,6 @@ rich safetensors scipy tb_nightly -tensordict toml torchdiffeq voluptuous @@ -41,34 +41,31 @@ scikit-image basicsr fasteners dctorch -pymatting -matplotlib -peft -orjson -httpx==0.24.1 compel==2.0.2 torchsde==0.2.6 -clip-interrogator==0.6.0 antlr4-python3-runtime==4.9.3 requests==2.31.0 tqdm==4.66.1 -accelerate==0.26.1 -opencv-contrib-python-headless==4.8.1.78 -diffusers==0.25.1 +accelerate==0.20.3 +opencv-python-headless==4.7.0.72 +diffusers==0.21.4 einops==0.4.1 gradio==3.43.2 -huggingface_hub==0.20.3 +huggingface_hub==0.17.1 numexpr==2.8.4 -numpy==1.26.2 -numba==0.58.1 +numpy==1.24.4 +numba==0.57.1 pandas==1.5.3 protobuf==3.20.3 pytorch_lightning==1.9.4 -tokenizers==0.15.1 -transformers==4.37.1 +transformers==4.32.1 tomesd==0.1.3 -urllib3==1.26.18 -Pillow==10.2.0 -timm==0.9.12 +urllib3==1.26.15 +Pillow==9.5.0 +timm==0.9.7 pydantic==1.10.13 -typing-extensions==4.9.0 +typing-extensions==4.8.0 + +torch==1.13.1 +torchvision==0.14.1 +torch-directml==0.1.13.1.dev230413