From cef798bf5708f8510c010271cd86932907fd743e Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Tue, 7 Nov 2023 21:48:27 +0900 Subject: [PATCH] implement olive img2img --- modules/onnx.py | 194 ++++++++++++++++++++++++++++++++++++++++++- modules/sd_models.py | 9 +- 2 files changed, 192 insertions(+), 11 deletions(-) diff --git a/modules/onnx.py b/modules/onnx.py index 0b05da97c..87f169f08 100644 --- a/modules/onnx.py +++ b/modules/onnx.py @@ -1,7 +1,9 @@ import os +import PIL import json import torch import shutil +import inspect import importlib import numpy as np import onnxruntime as ort @@ -10,6 +12,7 @@ import optimum.onnxruntime from enum import Enum from abc import ABCMeta from typing import Union, Optional, Callable, Type, List, Any, Dict +from diffusers.image_processor import VaeImageProcessor from installer import log from modules import shared, olive from modules.paths import sd_configs_path @@ -22,7 +25,8 @@ class ExecutionProvider(str, Enum): ROCm = "ROCMExecutionProvider" OpenVINO = "OpenVINOExecutionProvider" -submodels = ("text_encoder", "unet", "vae_encoder", "vae_decoder",) +submodels_sd = ("text_encoder", "unet", "vae_encoder", "vae_decoder",) +submodels_sdxl = ("text_encoder", "text_encoder_2", "unet", "vae_encoder", "vae_decoder",) available_execution_providers: List[ExecutionProvider] = ort.get_available_providers() EP_TO_NAME = { @@ -162,8 +166,6 @@ class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline, OnnxPip 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 @@ -289,9 +291,12 @@ class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline, OnnxPip diffusers.OnnxStableDiffusionPipeline = OnnxStableDiffusionPipeline +diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion"] = diffusers.OnnxStableDiffusionPipeline class OnnxStableDiffusionImg2ImgPipeline(diffusers.OnnxStableDiffusionImg2ImgPipeline, OnnxPipelineBase): + image_processor: VaeImageProcessor + def __init__( self, vae_encoder: diffusers.OnnxRuntimeModel, @@ -305,6 +310,7 @@ class OnnxStableDiffusionImg2ImgPipeline(diffusers.OnnxStableDiffusionImg2ImgPip requires_safety_checker: bool = True ): super().__init__(vae_encoder, vae_decoder, text_encoder, tokenizer, unet, scheduler, safety_checker, feature_extractor, requires_safety_checker) + self.image_processor = VaeImageProcessor(vae_scale_factor=64) @staticmethod def from_pretrained(pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs): @@ -340,10 +346,179 @@ class OnnxStableDiffusionImg2ImgPipeline(diffusers.OnnxStableDiffusionImg2ImgPip pass return OnnxStableDiffusionImg2ImgPipeline(**init_kwargs) + def __call__( + self, + prompt: Union[str, List[str]], + image: Union[np.ndarray, PIL.Image.Image] = None, + strength: float = 0.8, + 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, + 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, + ): + # check inputs. Raise error if not correct + self.check_inputs(prompt, 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] + + if strength < 0 or strength > 1: + raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") + + # set timesteps + self.scheduler.set_timesteps(num_inference_steps) + + image = self.image_processor.preprocess(image).cpu().numpy() + + # 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, + ) + + latents_dtype = prompt_embeds.dtype + image = image.astype(latents_dtype) + # encode the init image into latents and scale the latents + init_latents = self.vae_encoder(sample=image)[0] + init_latents = 0.18215 * init_latents + + if isinstance(prompt, str): + prompt = [prompt] + + init_latents = np.concatenate([init_latents] * num_images_per_prompt, axis=0) + + # get the original timestep using init_timestep + offset = self.scheduler.config.get("steps_offset", 0) + init_timestep = int(num_inference_steps * strength) + offset + init_timestep = min(init_timestep, num_inference_steps) + + 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) + init_latents = self.scheduler.add_noise( + torch.from_numpy(init_latents), torch.from_numpy(noise), torch.from_numpy(timesteps) + ) + init_latents = init_latents.numpy() + + # 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 + + latents = init_latents + + t_start = max(num_inference_steps - init_timestep + offset, 0) + timesteps = self.scheduler.timesteps[t_start:].numpy() + + 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(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)[ + 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 + + # skip postprocess + + if not return_dict: + return (image, has_nsfw_concept) + + return diffusers.pipelines.stable_diffusion.StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) + OnnxStableDiffusionImg2ImgPipeline.__module__ = 'diffusers' OnnxStableDiffusionImg2ImgPipeline.__name__ = 'OnnxStableDiffusionImg2ImgPipeline' diffusers.OnnxStableDiffusionImg2ImgPipeline = OnnxStableDiffusionImg2ImgPipeline +diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion"] = diffusers.OnnxStableDiffusionImg2ImgPipeline class OnnxStableDiffusionXLPipeline(optimum.onnxruntime.ORTStableDiffusionXLPipeline, OnnxPipelineBase): @@ -369,6 +544,7 @@ class OnnxStableDiffusionXLPipeline(optimum.onnxruntime.ORTStableDiffusionXLPipe OnnxStableDiffusionXLPipeline.__module__ = 'optimum.onnxruntime.modeling_diffusion' OnnxStableDiffusionXLPipeline.__name__ = 'ORTStableDiffusionXLPipeline' diffusers.OnnxStableDiffusionXLPipeline = OnnxStableDiffusionXLPipeline +diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion-xl"] = diffusers.OnnxStableDiffusionXLPipeline class OnnxStableDiffusionXLImg2ImgPipeline(optimum.onnxruntime.ORTStableDiffusionXLImg2ImgPipeline, OnnxPipelineBase): @@ -394,6 +570,7 @@ class OnnxStableDiffusionXLImg2ImgPipeline(optimum.onnxruntime.ORTStableDiffusio OnnxStableDiffusionXLImg2ImgPipeline.__module__ = 'optimum.onnxruntime.modeling_diffusion' OnnxStableDiffusionXLImg2ImgPipeline.__name__ = 'ORTStableDiffusionXLImg2ImgPipeline' diffusers.OnnxStableDiffusionXLImg2ImgPipeline = OnnxStableDiffusionXLImg2ImgPipeline +diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion-xl"] = diffusers.OnnxStableDiffusionXLImg2ImgPipeline class OnnxAutoPipelineBase(OnnxPipelineBase): @@ -408,10 +585,19 @@ class OnnxAutoPipelineBase(OnnxPipelineBase): self.pipeline = pipeline del pipeline + @property + def scheduler(self): + return self.pipeline.scheduler + + @scheduler.setter + def scheduler(self, scheduler): + self.pipeline.scheduler = scheduler + def derive_properties(self, pipeline: OnnxPipelineBase): pipeline.sd_model_hash = self.sd_model_hash pipeline.sd_checkpoint_info = self.sd_checkpoint_info pipeline.sd_model_checkpoint = self.sd_model_checkpoint + pipeline.scheduler = self.scheduler return pipeline def to(self, *args, **kwargs): @@ -458,6 +644,7 @@ class OnnxAutoPipelineBase(OnnxPipelineBase): shared.opts.onnx_temp_dir, out_dir, ignore=shutil.ignore_patterns("weights.pb", "*.onnx", "*.safetensors", "*.ckpt") ) + submodels = submodels_sdxl if olive.is_sdxl else submodels_sd converted_model_paths = {} for submodel in submodels: @@ -565,6 +752,7 @@ class OnnxAutoPipelineBase(OnnxPipelineBase): in_dir, out_dir, ignore=shutil.ignore_patterns("weights.pb", "*.onnx", "*.safetensors", "*.ckpt") ) + submodels = submodels_sdxl if olive.is_sdxl else submodels_sd optimized_model_paths = {} for submodel in submodels: diff --git a/modules/sd_models.py b/modules/sd_models.py index 64664fcab..9823a5b8d 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -790,14 +790,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No shared.log.debug(f'Diffusers loading: path="{checkpoint_info.path}"') pipeline, model_type = detect_pipeline(checkpoint_info.path, op) - if 'ONNX' in shared.opts.diffusers_pipeline: - from modules.onnx import OnnxAutoPipelineForText2Image - if os.path.isdir(checkpoint_info.path): - sd_model = OnnxAutoPipelineForText2Image.from_pretrained(checkpoint_info.path) - else: - sd_model = OnnxAutoPipelineForText2Image.from_single_file(checkpoint_info.path) - - if sd_model is None and os.path.isdir(checkpoint_info.path): + if os.path.isdir(checkpoint_info.path): err1 = None err2 = None err3 = None