from typing import Any
from diffusers import (
DiffusionPipeline,
AutoencoderTiny,
LCMScheduler,
UNet2DConditionModel,
)
from os import path
import torch
from backend.models.lcmdiffusion_setting import LCMDiffusionSetting
import numpy as np
from constants import (
DEVICE,
LCM_DEFAULT_MODEL,
TAESD_MODEL,
TAESDXL_MODEL,
TAESD_MODEL_OPENVINO,
)
from huggingface_hub import model_info
from backend.models.lcmdiffusion_setting import LCMLora
from backend.device import is_openvino_device
if is_openvino_device():
from huggingface_hub import snapshot_download
from optimum.intel.openvino.modeling_diffusion import OVModelVaeDecoder, OVBaseModel
# from optimum.intel.openvino.modeling_diffusion import OVStableDiffusionPipeline
from backend.lcmdiffusion.pipelines.openvino.lcm_ov_pipeline import (
OVStableDiffusionPipeline,
)
from backend.lcmdiffusion.pipelines.openvino.lcm_scheduler import (
LCMScheduler as OpenVinoLCMscheduler,
)
class CustomOVModelVaeDecoder(OVModelVaeDecoder):
def __init__(
self,
model,
parent_model,
ov_config=None,
model_dir=None,
):
super(OVModelVaeDecoder, self).__init__(
model,
parent_model,
ov_config,
"vae_decoder",
model_dir,
)
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.torch_data_type = (
torch.float32 if is_openvino_device() or DEVICE == "mps" else torch.float16
)
print(f"Torch datatype : {self.torch_data_type}")
def _get_lcm_pipeline(
self,
lcm_model_id: str,
base_model_id: str,
use_local_model: bool,
):
pipeline = None
unet = UNet2DConditionModel.from_pretrained(
lcm_model_id,
torch_dtype=torch.float32,
local_files_only=use_local_model
# resume_download=True,
)
pipeline = DiffusionPipeline.from_pretrained(
base_model_id,
unet=unet,
torch_dtype=torch.float32,
local_files_only=use_local_model
# resume_download=True,
)
pipeline.scheduler = LCMScheduler.from_config(pipeline.scheduler.config)
return pipeline
def get_tiny_decoder_vae_model(self) -> str:
pipeline_class = self.pipeline.__class__.__name__
print(f"Pipeline class : {pipeline_class}")
if (
pipeline_class == "LatentConsistencyModelPipeline"
or pipeline_class == "StableDiffusionPipeline"
):
return TAESD_MODEL
elif pipeline_class == "StableDiffusionXLPipeline":
return TAESDXL_MODEL
elif pipeline_class == "OVStableDiffusionPipeline":
return TAESD_MODEL_OPENVINO
def _get_lcm_model_pipeline(
self,
model_id: str,
use_local_model,
):
pipeline = None
if model_id == LCM_DEFAULT_MODEL:
pipeline = DiffusionPipeline.from_pretrained(
model_id,
local_files_only=use_local_model,
)
elif model_id == "latent-consistency/lcm-sdxl":
pipeline = self._get_lcm_pipeline(
model_id,
"stabilityai/stable-diffusion-xl-base-1.0",
use_local_model,
)
elif model_id == "latent-consistency/lcm-ssd-1b":
pipeline = self._get_lcm_pipeline(
model_id,
"segmind/SSD-1B",
use_local_model,
)
return pipeline
def _get_lcm_lora_pipeline(
self,
base_model_id: str,
lcm_lora_id: str,
use_local_model: bool,
):
pipeline = DiffusionPipeline.from_pretrained(
base_model_id,
torch_dtype=self.torch_data_type,
local_files_only=use_local_model,
)
pipeline.load_lora_weights(
lcm_lora_id,
local_files_only=use_local_model,
)
pipeline.scheduler = LCMScheduler.from_config(pipeline.scheduler.config)
pipeline.fuse_lora()
pipeline.unet.to(memory_format=torch.channels_last)
return pipeline
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 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 init(
self,
model_id: str,
use_openvino: bool = False,
device: str = "cpu",
use_local_model: bool = False,
use_tiny_auto_encoder: bool = False,
use_lora: bool = False,
lcm_lora: LCMLora = LCMLora(),
) -> None:
self.device = device
self.use_openvino = use_openvino
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
):
if self.use_openvino and is_openvino_device():
if self.pipeline:
del self.pipeline
self.pipeline = None
self.pipeline = OVStableDiffusionPipeline.from_pretrained(
model_id,
local_files_only=use_local_model,
ov_config={"CACHE_DIR": ""},
device=DEVICE.upper(),
)
if use_tiny_auto_encoder:
print("Using Tiny Auto Encoder (OpenVINO)")
taesd_dir = snapshot_download(
repo_id=self.get_tiny_decoder_vae_model(),
local_files_only=use_local_model,
)
self.pipeline.vae_decoder = CustomOVModelVaeDecoder(
model=OVBaseModel.load_model(
f"{taesd_dir}/vae_decoder/openvino_model.xml"
),
parent_model=self.pipeline,
model_dir=taesd_dir,
)
else:
if self.pipeline:
del self.pipeline
self.pipeline = None
if use_lora:
print("Init LCM-LoRA pipeline")
self.pipeline = self._get_lcm_lora_pipeline(
lcm_lora.base_model_id,
lcm_lora.lcm_lora_id,
use_local_model,
)
else:
print("Init LCM Model pipeline")
self.pipeline = self._get_lcm_model_pipeline(
model_id,
use_local_model,
)
if use_tiny_auto_encoder:
vae_model = self.get_tiny_decoder_vae_model()
print(f"Using Tiny Auto Encoder {vae_model}")
self.pipeline.vae = AutoencoderTiny.from_pretrained(
vae_model,
torch_dtype=torch.float32,
local_files_only=use_local_model,
)
self._pipeline_to_device()
self.previous_model_id = 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
print(f"Model :{model_id}")
print(f"Pipeline : {self.pipeline}")
self.pipeline.scheduler = LCMScheduler.from_config(
self.pipeline.scheduler.config,
beta_start=0.001,
beta_end=0.01,
)
if use_lora:
self._add_freeu()
def generate(
self,
lcm_diffusion_setting: LCMDiffusionSetting,
reshape: bool = False,
) -> Any:
guidance_scale = lcm_diffusion_setting.guidance_scale
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)
if lcm_diffusion_setting.use_openvino and is_openvino_device():
print("Using OpenVINO")
if reshape:
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 not lcm_diffusion_setting.use_safety_checker:
self.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:
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
else:
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
return result_images