| from typing import Any | |
| from diffusers import LCMScheduler | |
| import torch | |
| from backend.models.lcmdiffusion_setting import LCMDiffusionSetting | |
| import numpy as np | |
| from constants import DEVICE | |
| from backend.models.lcmdiffusion_setting import LCMLora | |
| from backend.device import is_openvino_device | |
| from backend.openvino.pipelines import ( | |
| get_ov_text_to_image_pipeline, | |
| ov_load_taesd, | |
| get_ov_image_to_image_pipeline, | |
| ) | |
| from backend.pipelines.lcm import ( | |
| get_lcm_model_pipeline, | |
| load_taesd, | |
| get_image_to_image_pipeline, | |
| ) | |
| from backend.pipelines.lcm_lora import get_lcm_lora_pipeline | |
| from backend.models.lcmdiffusion_setting import DiffusionTask | |
| from image_ops import resize_pil_image | |
| from math import ceil | |
| class LCMTextToImage: | |
| def __init__( | |
| self, | |
| device: str = "cpu", | |
| ) -> None: | |
| self.pipeline = None | |
| self.use_openvino = False | |
| self.device = "" | |
| self.previous_model_id = None | |
| self.previous_use_tae_sd = False | |
| self.previous_use_lcm_lora = False | |
| self.previous_ov_model_id = "" | |
| self.previous_safety_checker = False | |
| self.previous_use_openvino = False | |
| self.img_to_img_pipeline = None | |
| self.is_openvino_init = False | |
| self.torch_data_type = ( | |
| torch.float32 if is_openvino_device() or DEVICE == "mps" else torch.float16 | |
| ) | |
| print(f"Torch datatype : {self.torch_data_type}") | |
| def _pipeline_to_device(self): | |
| print(f"Pipeline device : {DEVICE}") | |
| print(f"Pipeline dtype : {self.torch_data_type}") | |
| self.pipeline.to( | |
| torch_device=DEVICE, | |
| torch_dtype=self.torch_data_type, | |
| ) | |
| def _add_freeu(self): | |
| pipeline_class = self.pipeline.__class__.__name__ | |
| if isinstance(self.pipeline.scheduler, LCMScheduler): | |
| if pipeline_class == "StableDiffusionPipeline": | |
| print("Add FreeU - SD") | |
| self.pipeline.enable_freeu( | |
| s1=0.9, | |
| s2=0.2, | |
| b1=1.2, | |
| b2=1.4, | |
| ) | |
| elif pipeline_class == "StableDiffusionXLPipeline": | |
| print("Add FreeU - SDXL") | |
| self.pipeline.enable_freeu( | |
| s1=0.6, | |
| s2=0.4, | |
| b1=1.1, | |
| b2=1.2, | |
| ) | |
| def _update_lcm_scheduler_params(self): | |
| if isinstance(self.pipeline.scheduler, LCMScheduler): | |
| self.pipeline.scheduler = LCMScheduler.from_config( | |
| self.pipeline.scheduler.config, | |
| beta_start=0.001, | |
| beta_end=0.01, | |
| ) | |
| def init( | |
| self, | |
| device: str = "cpu", | |
| lcm_diffusion_setting: LCMDiffusionSetting = LCMDiffusionSetting(), | |
| ) -> None: | |
| self.device = device | |
| self.use_openvino = lcm_diffusion_setting.use_openvino | |
| model_id = lcm_diffusion_setting.lcm_model_id | |
| use_local_model = lcm_diffusion_setting.use_offline_model | |
| use_tiny_auto_encoder = lcm_diffusion_setting.use_tiny_auto_encoder | |
| use_lora = lcm_diffusion_setting.use_lcm_lora | |
| lcm_lora: LCMLora = lcm_diffusion_setting.lcm_lora | |
| ov_model_id = lcm_diffusion_setting.openvino_lcm_model_id | |
| if lcm_diffusion_setting.diffusion_task == DiffusionTask.image_to_image.value: | |
| lcm_diffusion_setting.init_image = resize_pil_image( | |
| lcm_diffusion_setting.init_image, | |
| lcm_diffusion_setting.image_width, | |
| lcm_diffusion_setting.image_height, | |
| ) | |
| if ( | |
| self.pipeline is None | |
| or self.previous_model_id != model_id | |
| or self.previous_use_tae_sd != use_tiny_auto_encoder | |
| or self.previous_lcm_lora_base_id != lcm_lora.base_model_id | |
| or self.previous_lcm_lora_id != lcm_lora.lcm_lora_id | |
| or self.previous_use_lcm_lora != use_lora | |
| or self.previous_ov_model_id != ov_model_id | |
| or self.previous_safety_checker != lcm_diffusion_setting.use_safety_checker | |
| or self.previous_use_openvino != lcm_diffusion_setting.use_openvino | |
| ): | |
| if self.use_openvino and is_openvino_device(): | |
| if self.pipeline: | |
| del self.pipeline | |
| self.pipeline = None | |
| self.is_openvino_init = True | |
| if ( | |
| lcm_diffusion_setting.diffusion_task | |
| == DiffusionTask.text_to_image.value | |
| ): | |
| print(f"***** Init Text to image (OpenVINO) - {ov_model_id} *****") | |
| self.pipeline = get_ov_text_to_image_pipeline( | |
| ov_model_id, | |
| use_local_model, | |
| ) | |
| elif ( | |
| lcm_diffusion_setting.diffusion_task | |
| == DiffusionTask.image_to_image.value | |
| ): | |
| print(f"***** Image to image (OpenVINO) - {ov_model_id} *****") | |
| self.pipeline = get_ov_image_to_image_pipeline( | |
| ov_model_id, | |
| use_local_model, | |
| ) | |
| else: | |
| if self.pipeline: | |
| del self.pipeline | |
| self.pipeline = None | |
| if self.img_to_img_pipeline: | |
| del self.img_to_img_pipeline | |
| self.img_to_img_pipeline = None | |
| if use_lora: | |
| print( | |
| f"***** Init LCM-LoRA pipeline - {lcm_lora.base_model_id} *****" | |
| ) | |
| self.pipeline = get_lcm_lora_pipeline( | |
| lcm_lora.base_model_id, | |
| lcm_lora.lcm_lora_id, | |
| use_local_model, | |
| torch_data_type=self.torch_data_type, | |
| ) | |
| else: | |
| print(f"***** Init LCM Model pipeline - {model_id} *****") | |
| self.pipeline = get_lcm_model_pipeline( | |
| model_id, | |
| use_local_model, | |
| ) | |
| if ( | |
| lcm_diffusion_setting.diffusion_task | |
| == DiffusionTask.image_to_image.value | |
| ): | |
| self.img_to_img_pipeline = get_image_to_image_pipeline( | |
| self.pipeline | |
| ) | |
| self._pipeline_to_device() | |
| if use_tiny_auto_encoder: | |
| if self.use_openvino and is_openvino_device(): | |
| print("Using Tiny Auto Encoder (OpenVINO)") | |
| ov_load_taesd( | |
| self.pipeline, | |
| use_local_model, | |
| ) | |
| else: | |
| print("Using Tiny Auto Encoder") | |
| if ( | |
| lcm_diffusion_setting.diffusion_task | |
| == DiffusionTask.text_to_image.value | |
| ): | |
| load_taesd( | |
| self.pipeline, | |
| use_local_model, | |
| self.torch_data_type, | |
| ) | |
| elif ( | |
| lcm_diffusion_setting.diffusion_task | |
| == DiffusionTask.image_to_image.value | |
| ): | |
| load_taesd( | |
| self.img_to_img_pipeline, | |
| use_local_model, | |
| self.torch_data_type, | |
| ) | |
| if ( | |
| lcm_diffusion_setting.diffusion_task | |
| == DiffusionTask.image_to_image.value | |
| and lcm_diffusion_setting.use_openvino | |
| ): | |
| self.pipeline.scheduler = LCMScheduler.from_config( | |
| self.pipeline.scheduler.config, | |
| ) | |
| else: | |
| self._update_lcm_scheduler_params() | |
| if use_lora: | |
| self._add_freeu() | |
| self.previous_model_id = model_id | |
| self.previous_ov_model_id = ov_model_id | |
| self.previous_use_tae_sd = use_tiny_auto_encoder | |
| self.previous_lcm_lora_base_id = lcm_lora.base_model_id | |
| self.previous_lcm_lora_id = lcm_lora.lcm_lora_id | |
| self.previous_use_lcm_lora = use_lora | |
| self.previous_safety_checker = lcm_diffusion_setting.use_safety_checker | |
| self.previous_use_openvino = lcm_diffusion_setting.use_openvino | |
| if ( | |
| lcm_diffusion_setting.diffusion_task | |
| == DiffusionTask.text_to_image.value | |
| ): | |
| print(f"Pipeline : {self.pipeline}") | |
| elif ( | |
| lcm_diffusion_setting.diffusion_task | |
| == DiffusionTask.image_to_image.value | |
| ): | |
| if self.use_openvino and is_openvino_device(): | |
| print(f"Pipeline : {self.pipeline}") | |
| else: | |
| print(f"Pipeline : {self.img_to_img_pipeline}") | |
| def generate( | |
| self, | |
| lcm_diffusion_setting: LCMDiffusionSetting, | |
| reshape: bool = False, | |
| ) -> Any: | |
| guidance_scale = lcm_diffusion_setting.guidance_scale | |
| img_to_img_inference_steps = lcm_diffusion_setting.inference_steps | |
| check_step_value = int( | |
| lcm_diffusion_setting.inference_steps * lcm_diffusion_setting.strength | |
| ) | |
| if ( | |
| lcm_diffusion_setting.diffusion_task == DiffusionTask.image_to_image.value | |
| and check_step_value < 1 | |
| ): | |
| img_to_img_inference_steps = ceil(1 / lcm_diffusion_setting.strength) | |
| print( | |
| f"Strength: {lcm_diffusion_setting.strength},{img_to_img_inference_steps}" | |
| ) | |
| if lcm_diffusion_setting.use_seed: | |
| cur_seed = lcm_diffusion_setting.seed | |
| if self.use_openvino: | |
| np.random.seed(cur_seed) | |
| else: | |
| torch.manual_seed(cur_seed) | |
| is_openvino_pipe = lcm_diffusion_setting.use_openvino and is_openvino_device() | |
| if is_openvino_pipe: | |
| print("Using OpenVINO") | |
| if reshape and not self.is_openvino_init: | |
| print("Reshape and compile") | |
| self.pipeline.reshape( | |
| batch_size=-1, | |
| height=lcm_diffusion_setting.image_height, | |
| width=lcm_diffusion_setting.image_width, | |
| num_images_per_prompt=lcm_diffusion_setting.number_of_images, | |
| ) | |
| self.pipeline.compile() | |
| if self.is_openvino_init: | |
| self.is_openvino_init = False | |
| if not lcm_diffusion_setting.use_safety_checker: | |
| self.pipeline.safety_checker = None | |
| if ( | |
| lcm_diffusion_setting.diffusion_task | |
| == DiffusionTask.image_to_image.value | |
| and not is_openvino_pipe | |
| ): | |
| self.img_to_img_pipeline.safety_checker = None | |
| if ( | |
| not lcm_diffusion_setting.use_lcm_lora | |
| and not lcm_diffusion_setting.use_openvino | |
| and lcm_diffusion_setting.guidance_scale != 1.0 | |
| ): | |
| print("Not using LCM-LoRA so setting guidance_scale 1.0") | |
| guidance_scale = 1.0 | |
| if lcm_diffusion_setting.use_openvino: | |
| if ( | |
| lcm_diffusion_setting.diffusion_task | |
| == DiffusionTask.text_to_image.value | |
| ): | |
| result_images = self.pipeline( | |
| prompt=lcm_diffusion_setting.prompt, | |
| negative_prompt=lcm_diffusion_setting.negative_prompt, | |
| num_inference_steps=lcm_diffusion_setting.inference_steps, | |
| guidance_scale=guidance_scale, | |
| width=lcm_diffusion_setting.image_width, | |
| height=lcm_diffusion_setting.image_height, | |
| num_images_per_prompt=lcm_diffusion_setting.number_of_images, | |
| ).images | |
| elif ( | |
| lcm_diffusion_setting.diffusion_task | |
| == DiffusionTask.image_to_image.value | |
| ): | |
| result_images = self.pipeline( | |
| image=lcm_diffusion_setting.init_image, | |
| strength=lcm_diffusion_setting.strength, | |
| prompt=lcm_diffusion_setting.prompt, | |
| negative_prompt=lcm_diffusion_setting.negative_prompt, | |
| num_inference_steps=img_to_img_inference_steps * 3, | |
| guidance_scale=guidance_scale, | |
| num_images_per_prompt=lcm_diffusion_setting.number_of_images, | |
| ).images | |
| else: | |
| if ( | |
| lcm_diffusion_setting.diffusion_task | |
| == DiffusionTask.text_to_image.value | |
| ): | |
| result_images = self.pipeline( | |
| prompt=lcm_diffusion_setting.prompt, | |
| negative_prompt=lcm_diffusion_setting.negative_prompt, | |
| num_inference_steps=lcm_diffusion_setting.inference_steps, | |
| guidance_scale=guidance_scale, | |
| width=lcm_diffusion_setting.image_width, | |
| height=lcm_diffusion_setting.image_height, | |
| num_images_per_prompt=lcm_diffusion_setting.number_of_images, | |
| ).images | |
| elif ( | |
| lcm_diffusion_setting.diffusion_task | |
| == DiffusionTask.image_to_image.value | |
| ): | |
| result_images = self.img_to_img_pipeline( | |
| image=lcm_diffusion_setting.init_image, | |
| strength=lcm_diffusion_setting.strength, | |
| prompt=lcm_diffusion_setting.prompt, | |
| negative_prompt=lcm_diffusion_setting.negative_prompt, | |
| num_inference_steps=img_to_img_inference_steps, | |
| guidance_scale=guidance_scale, | |
| width=lcm_diffusion_setting.image_width, | |
| height=lcm_diffusion_setting.image_height, | |
| num_images_per_prompt=lcm_diffusion_setting.number_of_images, | |
| ).images | |
| return result_images | |