import gc
from math import ceil
from typing import Any, List
import random
import numpy as np
import torch
import logging
from backend.device import is_openvino_device
from backend.lora import load_lora_weight
from backend.controlnet import (
load_controlnet_adapters,
update_controlnet_arguments,
)
from backend.models.lcmdiffusion_setting import (
DiffusionTask,
LCMDiffusionSetting,
LCMLora,
)
from backend.openvino.pipelines import (
get_ov_image_to_image_pipeline,
get_ov_text_to_image_pipeline,
ov_load_taesd,
)
from backend.pipelines.lcm import (
get_image_to_image_pipeline,
get_lcm_model_pipeline,
load_taesd,
)
from backend.pipelines.lcm_lora import get_lcm_lora_pipeline
from constants import DEVICE, GGUF_THREADS
from diffusers import LCMScheduler
from image_ops import resize_pil_image
from backend.openvino.flux_pipeline import get_flux_pipeline
from backend.openvino.ov_hc_stablediffusion_pipeline import OvHcLatentConsistency
from backend.gguf.gguf_diffusion import (
GGUFDiffusion,
ModelConfig,
Txt2ImgConfig,
SampleMethod,
)
from paths import get_app_path
from pprint import pprint
try:
# support for token merging; keeping it optional for now
import tomesd
except ImportError:
print("tomesd library unavailable; disabling token merging support")
tomesd = None
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_token_merging = 0.0
self.previous_safety_checker = False
self.previous_use_openvino = False
self.img_to_img_pipeline = None
self.is_openvino_init = False
self.previous_lora = None
self.task_type = DiffusionTask.text_to_image
self.previous_use_gguf_model = False
self.previous_gguf_model = None
self.torch_data_type = (
torch.float32 if is_openvino_device() or DEVICE == "mps" else torch.float16
)
self.ov_model_id = None
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 _enable_vae_tiling(self):
self.pipeline.vae.enable_tiling()
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 _is_hetero_pipeline(self) -> bool:
return "square" in self.ov_model_id.lower()
def _load_ov_hetero_pipeline(self):
print("Loading Heterogeneous Compute pipeline")
if DEVICE.upper()=="NPU":
device = ["NPU", "NPU", "NPU"]
self.pipeline = OvHcLatentConsistency(self.ov_model_id,device)
else:
self.pipeline = OvHcLatentConsistency(self.ov_model_id)
def _generate_images_hetero_compute(
self,
lcm_diffusion_setting: LCMDiffusionSetting,
):
print("Using OpenVINO ")
if lcm_diffusion_setting.diffusion_task == DiffusionTask.text_to_image.value:
return [
self.pipeline.generate(
prompt=lcm_diffusion_setting.prompt,
neg_prompt=lcm_diffusion_setting.negative_prompt,
init_image=None,
strength=1.0,
num_inference_steps=lcm_diffusion_setting.inference_steps,
)
]
else:
return [
self.pipeline.generate(
prompt=lcm_diffusion_setting.prompt,
neg_prompt=lcm_diffusion_setting.negative_prompt,
init_image=lcm_diffusion_setting.init_image,
strength=lcm_diffusion_setting.strength,
num_inference_steps=lcm_diffusion_setting.inference_steps,
)
]
def _is_valid_mode(
self,
modes: List,
) -> bool:
return modes.count(True) == 1 or modes.count(False) == 3
def _validate_mode(
self,
modes: List,
) -> None:
if not self._is_valid_mode(modes):
raise ValueError("Invalid mode,delete configs/settings.yaml and retry!")
def init(
self,
device: str = "cpu",
lcm_diffusion_setting: LCMDiffusionSetting = LCMDiffusionSetting(),
) -> None:
# Mode validation either LCM LoRA or OpenVINO or GGUF
modes = [
lcm_diffusion_setting.use_gguf_model,
lcm_diffusion_setting.use_openvino,
lcm_diffusion_setting.use_lcm_lora,
]
self._validate_mode(modes)
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
token_merging = lcm_diffusion_setting.token_merging
self.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 != self.ov_model_id
or self.previous_token_merging != token_merging
or self.previous_safety_checker != lcm_diffusion_setting.use_safety_checker
or self.previous_use_openvino != lcm_diffusion_setting.use_openvino
or self.previous_use_gguf_model != lcm_diffusion_setting.use_gguf_model
or self.previous_gguf_model != lcm_diffusion_setting.gguf_model
or (
self.use_openvino
and (
self.previous_task_type != lcm_diffusion_setting.diffusion_task
or self.previous_lora != lcm_diffusion_setting.lora
)
)
or lcm_diffusion_setting.rebuild_pipeline
):
if self.use_openvino and is_openvino_device():
if self.pipeline:
del self.pipeline
self.pipeline = None
gc.collect()
self.is_openvino_init = True
if (
lcm_diffusion_setting.diffusion_task
== DiffusionTask.text_to_image.value
):
print(
f"***** Init Text to image (OpenVINO) - {self.ov_model_id} *****"
)
if "flux" in self.ov_model_id.lower():
print("Loading OpenVINO Flux pipeline")
self.pipeline = get_flux_pipeline(
self.ov_model_id,
lcm_diffusion_setting.use_tiny_auto_encoder,
)
elif self._is_hetero_pipeline():
self._load_ov_hetero_pipeline()
else:
self.pipeline = get_ov_text_to_image_pipeline(
self.ov_model_id,
use_local_model,
)
elif (
lcm_diffusion_setting.diffusion_task
== DiffusionTask.image_to_image.value
):
if not self.pipeline and self._is_hetero_pipeline():
self._load_ov_hetero_pipeline()
else:
print(
f"***** Image to image (OpenVINO) - {self.ov_model_id} *****"
)
self.pipeline = get_ov_image_to_image_pipeline(
self.ov_model_id,
use_local_model,
)
elif lcm_diffusion_setting.use_gguf_model:
model = lcm_diffusion_setting.gguf_model.diffusion_path
print(f"***** Init Text to image (GGUF) - {model} *****")
# if self.pipeline:
# self.pipeline.terminate()
# del self.pipeline
# self.pipeline = None
self._init_gguf_diffusion(lcm_diffusion_setting)
else:
if self.pipeline or self.img_to_img_pipeline:
self.pipeline = None
self.img_to_img_pipeline = None
gc.collect()
controlnet_args = load_controlnet_adapters(lcm_diffusion_setting)
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,
pipeline_args=controlnet_args,
)
else:
print(f"***** Init LCM Model pipeline - {model_id} *****")
self.pipeline = get_lcm_model_pipeline(
model_id,
use_local_model,
controlnet_args,
)
self.img_to_img_pipeline = get_image_to_image_pipeline(self.pipeline)
if tomesd and token_merging > 0.001:
print(f"***** Token Merging: {token_merging} *****")
tomesd.apply_patch(self.pipeline, ratio=token_merging)
tomesd.apply_patch(self.img_to_img_pipeline, ratio=token_merging)
if use_tiny_auto_encoder:
if self.use_openvino and is_openvino_device():
if self.pipeline.__class__.__name__ != "OVFluxPipeline":
print("Using Tiny Auto Encoder (OpenVINO)")
ov_load_taesd(
self.pipeline,
use_local_model,
)
else:
print("Using Tiny Auto Encoder")
load_taesd(
self.pipeline,
use_local_model,
self.torch_data_type,
)
load_taesd(
self.img_to_img_pipeline,
use_local_model,
self.torch_data_type,
)
if not self.use_openvino and not is_openvino_device():
self._pipeline_to_device()
if not self._is_hetero_pipeline():
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:
if not lcm_diffusion_setting.use_gguf_model:
self._update_lcm_scheduler_params()
if use_lora:
self._add_freeu()
self.previous_model_id = model_id
self.previous_ov_model_id = self.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_token_merging = lcm_diffusion_setting.token_merging
self.previous_safety_checker = lcm_diffusion_setting.use_safety_checker
self.previous_use_openvino = lcm_diffusion_setting.use_openvino
self.previous_task_type = lcm_diffusion_setting.diffusion_task
self.previous_lora = lcm_diffusion_setting.lora.model_copy(deep=True)
self.previous_use_gguf_model = lcm_diffusion_setting.use_gguf_model
self.previous_gguf_model = lcm_diffusion_setting.gguf_model.model_copy(
deep=True
)
lcm_diffusion_setting.rebuild_pipeline = False
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}")
if self.use_openvino:
if lcm_diffusion_setting.lora.enabled:
print("Warning: Lora models not supported on OpenVINO mode")
elif not lcm_diffusion_setting.use_gguf_model:
adapters = self.pipeline.get_active_adapters()
print(f"Active adapters : {adapters}")
def _get_timesteps(self):
time_steps = self.pipeline.scheduler.config.get("timesteps")
time_steps_value = [int(time_steps)] if time_steps else None
return time_steps_value
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}"
)
pipeline_extra_args = {}
if lcm_diffusion_setting.use_seed:
cur_seed = lcm_diffusion_setting.seed
# for multiple images with a fixed seed, use sequential seeds
seeds = [(cur_seed + i) for i in range(lcm_diffusion_setting.number_of_images)]
else:
seeds = [random.randint(0,999999999) for i in range(lcm_diffusion_setting.number_of_images)]
if self.use_openvino:
# no support for generators; try at least to ensure reproducible results for single images
np.random.seed(seeds[0])
if self._is_hetero_pipeline():
torch.manual_seed(seeds[0])
lcm_diffusion_setting.seed=seeds[0]
else:
pipeline_extra_args['generator'] = [
torch.Generator(device=self.device).manual_seed(s) for s in seeds]
is_openvino_pipe = lcm_diffusion_setting.use_openvino and is_openvino_device()
if is_openvino_pipe and not self._is_hetero_pipeline():
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 is_openvino_pipe and self._is_hetero_pipeline():
return self._generate_images_hetero_compute(lcm_diffusion_setting)
elif lcm_diffusion_setting.use_gguf_model:
return self._generate_images_gguf(lcm_diffusion_setting)
if lcm_diffusion_setting.clip_skip > 1:
# We follow the convention that "CLIP Skip == 2" means "skip
# the last layer", so "CLIP Skip == 1" means "no skipping"
pipeline_extra_args["clip_skip"] = lcm_diffusion_setting.clip_skip - 1
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
controlnet_args = update_controlnet_arguments(lcm_diffusion_setting)
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,
timesteps=self._get_timesteps(),
**pipeline_extra_args,
**controlnet_args,
).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,
**pipeline_extra_args,
**controlnet_args,
).images
for (i, seed) in enumerate(seeds):
result_images[i].info['image_seed'] = seed
return result_images
def _init_gguf_diffusion(
self,
lcm_diffusion_setting: LCMDiffusionSetting,
):
config = ModelConfig()
config.model_path = lcm_diffusion_setting.gguf_model.diffusion_path
config.diffusion_model_path = lcm_diffusion_setting.gguf_model.diffusion_path
config.clip_l_path = lcm_diffusion_setting.gguf_model.clip_path
config.t5xxl_path = lcm_diffusion_setting.gguf_model.t5xxl_path
config.vae_path = lcm_diffusion_setting.gguf_model.vae_path
config.n_threads = GGUF_THREADS
print(f"GGUF Threads : {GGUF_THREADS} ")
print("GGUF - Model config")
pprint(lcm_diffusion_setting.gguf_model.model_dump())
self.pipeline = GGUFDiffusion(
get_app_path(), # Place DLL in fastsdcpu folder
config,
True,
)
def _generate_images_gguf(
self,
lcm_diffusion_setting: LCMDiffusionSetting,
):
if lcm_diffusion_setting.diffusion_task == DiffusionTask.text_to_image.value:
t2iconfig = Txt2ImgConfig()
t2iconfig.prompt = lcm_diffusion_setting.prompt
t2iconfig.batch_count = lcm_diffusion_setting.number_of_images
t2iconfig.cfg_scale = lcm_diffusion_setting.guidance_scale
t2iconfig.height = lcm_diffusion_setting.image_height
t2iconfig.width = lcm_diffusion_setting.image_width
t2iconfig.sample_steps = lcm_diffusion_setting.inference_steps
t2iconfig.sample_method = SampleMethod.EULER
if lcm_diffusion_setting.use_seed:
t2iconfig.seed = lcm_diffusion_setting.seed
else:
t2iconfig.seed = -1
return self.pipeline.generate_text2mg(t2iconfig)