| 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 |
|
|
