"""
Copyright(C) 2022-2023 Intel Corporation
SPDX - License - Identifier: Apache - 2.0
"""
import inspect
from typing import Union, Optional, Any, List, Dict
import numpy as np
# openvino
from openvino.runtime import Core
# tokenizer
from transformers import CLIPTokenizer
import torch
import random
from diffusers import DiffusionPipeline
from diffusers.schedulers import (DDIMScheduler,
LMSDiscreteScheduler,
PNDMScheduler,
EulerDiscreteScheduler,
EulerAncestralDiscreteScheduler)
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils.torch_utils import randn_tensor
from diffusers.utils import PIL_INTERPOLATION
import cv2
import os
import sys
# for multithreading
import concurrent.futures
#For GIF
import PIL
from PIL import Image
import glob
import json
import time
def scale_fit_to_window(dst_width:int, dst_height:int, image_width:int, image_height:int):
"""
Preprocessing helper function for calculating image size for resize with peserving original aspect ratio
and fitting image to specific window size
Parameters:
dst_width (int): destination window width
dst_height (int): destination window height
image_width (int): source image width
image_height (int): source image height
Returns:
result_width (int): calculated width for resize
result_height (int): calculated height for resize
"""
im_scale = min(dst_height / image_height, dst_width / image_width)
return int(im_scale * image_width), int(im_scale * image_height)
def preprocess(image: PIL.Image.Image, ht=512, wt=512):
"""
Image preprocessing function. Takes image in PIL.Image format, resizes it to keep aspect ration and fits to model input window 512x512,
then converts it to np.ndarray and adds padding with zeros on right or bottom side of image (depends from aspect ratio), after that
converts data to float32 data type and change range of values from [0, 255] to [-1, 1], finally, converts data layout from planar NHWC to NCHW.
The function returns preprocessed input tensor and padding size, which can be used in postprocessing.
Parameters:
image (PIL.Image.Image): input image
Returns:
image (np.ndarray): preprocessed image tensor
meta (Dict): dictionary with preprocessing metadata info
"""
src_width, src_height = image.size
image = image.convert('RGB')
dst_width, dst_height = scale_fit_to_window(
wt, ht, src_width, src_height)
image = np.array(image.resize((dst_width, dst_height),
resample=PIL.Image.Resampling.LANCZOS))[None, :]
pad_width = wt - dst_width
pad_height = ht - dst_height
pad = ((0, 0), (0, pad_height), (0, pad_width), (0, 0))
image = np.pad(image, pad, mode="constant")
image = image.astype(np.float32) / 255.0
image = 2.0 * image - 1.0
image = image.transpose(0, 3, 1, 2)
return image, {"padding": pad, "src_width": src_width, "src_height": src_height}
def try_enable_npu_turbo(device, core):
import platform
if "windows" in platform.system().lower():
if "NPU" in device and "3720" not in core.get_property('NPU', 'DEVICE_ARCHITECTURE'):
try:
core.set_property(properties={'NPU_TURBO': 'YES'},device_name='NPU')
except:
print(f"Failed loading NPU_TURBO for device {device}. Skipping... ")
else:
print_npu_turbo_art()
else:
print(f"Skipping NPU_TURBO for device {device}")
elif "linux" in platform.system().lower():
if os.path.isfile('/sys/module/intel_vpu/parameters/test_mode'):
with open('/sys/module/intel_vpu/version', 'r') as f:
version = f.readline().split()[0]
if tuple(map(int, version.split('.'))) < tuple(map(int, '1.9.0'.split('.'))):
print(f"The driver intel_vpu-1.9.0 (or later) needs to be loaded for NPU Turbo (currently {version}). Skipping...")
else:
with open('/sys/module/intel_vpu/parameters/test_mode', 'r') as tm_file:
test_mode = int(tm_file.readline().split()[0])
if test_mode == 512:
print_npu_turbo_art()
else:
print("The driver >=intel_vpu-1.9.0 was must be loaded with "
"\"modprobe intel_vpu test_mode=512\" to enable NPU_TURBO "
f"(currently test_mode={test_mode}). Skipping...")
else:
print(f"The driver >=intel_vpu-1.9.0 must be loaded with \"modprobe intel_vpu test_mode=512\" to enable NPU_TURBO. Skipping...")
else:
print(f"This platform ({platform.system()}) does not support NPU Turbo")
def result(var):
return next(iter(var.values()))
class StableDiffusionEngineAdvanced(DiffusionPipeline):
def __init__(self, model="runwayml/stable-diffusion-v1-5",
tokenizer="openai/clip-vit-large-patch14",
device=["CPU", "CPU", "CPU", "CPU"]):
try:
self.tokenizer = CLIPTokenizer.from_pretrained(model, local_files_only=True)
except:
self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer)
self.tokenizer.save_pretrained(model)
self.core = Core()
self.core.set_property({'CACHE_DIR': os.path.join(model, 'cache')})
try_enable_npu_turbo(device, self.core)
print("Loading models... ")
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
futures = {
"unet_time_proj": executor.submit(self.core.compile_model, os.path.join(model, "unet_time_proj.xml"), device[0]),
"text": executor.submit(self.load_model, model, "text_encoder", device[0]),
"unet": executor.submit(self.load_model, model, "unet_int8", device[1]),
"unet_neg": executor.submit(self.load_model, model, "unet_int8", device[2]) if device[1] != device[2] else None,
"vae_decoder": executor.submit(self.load_model, model, "vae_decoder", device[3]),
"vae_encoder": executor.submit(self.load_model, model, "vae_encoder", device[3])
}
self.unet_time_proj = futures["unet_time_proj"].result()
self.text_encoder = futures["text"].result()
self.unet = futures["unet"].result()
self.unet_neg = futures["unet_neg"].result() if futures["unet_neg"] else self.unet
self.vae_decoder = futures["vae_decoder"].result()
self.vae_encoder = futures["vae_encoder"].result()
print("Text Device:", device[0])
print("unet Device:", device[1])
print("unet-neg Device:", device[2])
print("VAE Device:", device[3])
self._text_encoder_output = self.text_encoder.output(0)
self._vae_d_output = self.vae_decoder.output(0)
self._vae_e_output = self.vae_encoder.output(0) if self.vae_encoder else None
self.set_dimensions()
self.infer_request_neg = self.unet_neg.create_infer_request()
self.infer_request = self.unet.create_infer_request()
self.infer_request_time_proj = self.unet_time_proj.create_infer_request()
self.time_proj_constants = np.load(os.path.join(model, "time_proj_constants.npy"))
def load_model(self, model, model_name, device):
if "NPU" in device:
with open(os.path.join(model, f"{model_name}.blob"), "rb") as f:
return self.core.import_model(f.read(), device)
return self.core.compile_model(os.path.join(model, f"{model_name}.xml"), device)
def set_dimensions(self):
latent_shape = self.unet.input("latent_model_input").shape
if latent_shape[1] == 4:
self.height = latent_shape[2] * 8
self.width = latent_shape[3] * 8
else:
self.height = latent_shape[1] * 8
self.width = latent_shape[2] * 8
def __call__(
self,
prompt,
init_image = None,
negative_prompt=None,
scheduler=None,
strength = 0.5,
num_inference_steps = 32,
guidance_scale = 7.5,
eta = 0.0,
create_gif = False,
model = None,
callback = None,
callback_userdata = None
):
# extract condition
text_input = self.tokenizer(
prompt,
padding="max_length",
max_length=self.tokenizer.model_max_length,
truncation=True,
return_tensors="np",
)
text_embeddings = self.text_encoder(text_input.input_ids)[self._text_encoder_output]
# do classifier free guidance
do_classifier_free_guidance = guidance_scale > 1.0
if do_classifier_free_guidance:
if negative_prompt is None:
uncond_tokens = [""]
elif isinstance(negative_prompt, str):
uncond_tokens = [negative_prompt]
else:
uncond_tokens = negative_prompt
tokens_uncond = self.tokenizer(
uncond_tokens,
padding="max_length",
max_length=self.tokenizer.model_max_length, #truncation=True,
return_tensors="np"
)
uncond_embeddings = self.text_encoder(tokens_uncond.input_ids)[self._text_encoder_output]
text_embeddings = np.concatenate([uncond_embeddings, text_embeddings])
# set timesteps
accepts_offset = "offset" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
extra_set_kwargs = {}
if accepts_offset:
extra_set_kwargs["offset"] = 1
scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)
timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, scheduler)
latent_timestep = timesteps[:1]
# get the initial random noise unless the user supplied it
latents, meta = self.prepare_latents(init_image, latent_timestep, scheduler)
# 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(scheduler.step).parameters.keys())
extra_step_kwargs = {}
if accepts_eta:
extra_step_kwargs["eta"] = eta
if create_gif:
frames = []
for i, t in enumerate(self.progress_bar(timesteps)):
if callback:
callback(i, callback_userdata)
# expand the latents if we are doing classifier free guidance
noise_pred = []
latent_model_input = latents
latent_model_input = scheduler.scale_model_input(latent_model_input, t)
latent_model_input_neg = latent_model_input
if self.unet.input("latent_model_input").shape[1] != 4:
#print("In transpose")
try:
latent_model_input = latent_model_input.permute(0,2,3,1)
except:
latent_model_input = latent_model_input.transpose(0,2,3,1)
if self.unet_neg.input("latent_model_input").shape[1] != 4:
#print("In transpose")
try:
latent_model_input_neg = latent_model_input_neg.permute(0,2,3,1)
except:
latent_model_input_neg = latent_model_input_neg.transpose(0,2,3,1)
time_proj_constants_fp16 = np.float16(self.time_proj_constants)
t_scaled_fp16 = time_proj_constants_fp16 * np.float16(t)
cosine_t_fp16 = np.cos(t_scaled_fp16)
sine_t_fp16 = np.sin(t_scaled_fp16)
t_scaled = self.time_proj_constants * np.float32(t)
cosine_t = np.cos(t_scaled)
sine_t = np.sin(t_scaled)
time_proj_dict = {"sine_t" : np.float32(sine_t), "cosine_t" : np.float32(cosine_t)}
self.infer_request_time_proj.start_async(time_proj_dict)
self.infer_request_time_proj.wait()
time_proj = self.infer_request_time_proj.get_output_tensor(0).data.astype(np.float32)
input_tens_neg_dict = {"time_proj": np.float32(time_proj), "latent_model_input":latent_model_input_neg, "encoder_hidden_states": np.expand_dims(text_embeddings[0], axis=0)}
input_tens_dict = {"time_proj": np.float32(time_proj), "latent_model_input":latent_model_input, "encoder_hidden_states": np.expand_dims(text_embeddings[1], axis=0)}
self.infer_request_neg.start_async(input_tens_neg_dict)
self.infer_request.start_async(input_tens_dict)
self.infer_request_neg.wait()
self.infer_request.wait()
noise_pred_neg = self.infer_request_neg.get_output_tensor(0)
noise_pred_pos = self.infer_request.get_output_tensor(0)
noise_pred.append(noise_pred_neg.data.astype(np.float32))
noise_pred.append(noise_pred_pos.data.astype(np.float32))
# perform guidance
if do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred[0], noise_pred[1]
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents = scheduler.step(torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs)["prev_sample"].numpy()
if create_gif:
frames.append(latents)
if callback:
callback(num_inference_steps, callback_userdata)
# scale and decode the image latents with vae
latents = 1 / 0.18215 * latents
start = time.time()
image = self.vae_decoder(latents)[self._vae_d_output]
print("Decoder ended:",time.time() - start)
image = self.postprocess_image(image, meta)
if create_gif:
gif_folder=os.path.join(model,"../../../gif")
print("gif_folder:",gif_folder)
if not os.path.exists(gif_folder):
os.makedirs(gif_folder)
for i in range(0,len(frames)):
image = self.vae_decoder(frames[i]*(1/0.18215))[self._vae_d_output]
image = self.postprocess_image(image, meta)
output = gif_folder + "/" + str(i).zfill(3) +".png"
cv2.imwrite(output, image)
with open(os.path.join(gif_folder, "prompt.json"), "w") as file:
json.dump({"prompt": prompt}, file)
frames_image = [Image.open(image) for image in glob.glob(f"{gif_folder}/*.png")]
frame_one = frames_image[0]
gif_file=os.path.join(gif_folder,"stable_diffusion.gif")
frame_one.save(gif_file, format="GIF", append_images=frames_image, save_all=True, duration=100, loop=0)
return image
def prepare_latents(self, image:PIL.Image.Image = None, latent_timestep:torch.Tensor = None, scheduler = LMSDiscreteScheduler):
"""
Function for getting initial latents for starting generation
Parameters:
image (PIL.Image.Image, *optional*, None):
Input image for generation, if not provided randon noise will be used as starting point
latent_timestep (torch.Tensor, *optional*, None):
Predicted by scheduler initial step for image generation, required for latent image mixing with nosie
Returns:
latents (np.ndarray):
Image encoded in latent space
"""
latents_shape = (1, 4, self.height // 8, self.width // 8)
noise = np.random.randn(*latents_shape).astype(np.float32)
if image is None:
##print("Image is NONE")
# if we use LMSDiscreteScheduler, let's make sure latents are mulitplied by sigmas
if isinstance(scheduler, LMSDiscreteScheduler):
noise = noise * scheduler.sigmas[0].numpy()
return noise, {}
elif isinstance(scheduler, EulerDiscreteScheduler) or isinstance(scheduler,EulerAncestralDiscreteScheduler):
noise = noise * scheduler.sigmas.max().numpy()
return noise, {}
else:
return noise, {}
input_image, meta = preprocess(image,self.height,self.width)
moments = self.vae_encoder(input_image)[self._vae_e_output]
mean, logvar = np.split(moments, 2, axis=1)
std = np.exp(logvar * 0.5)
latents = (mean + std * np.random.randn(*mean.shape)) * 0.18215
latents = scheduler.add_noise(torch.from_numpy(latents), torch.from_numpy(noise), latent_timestep).numpy()
return latents, meta
def postprocess_image(self, image:np.ndarray, meta:Dict):
"""
Postprocessing for decoded image. Takes generated image decoded by VAE decoder, unpad it to initial image size (if required),
normalize and convert to [0, 255] pixels range. Optionally, convertes it from np.ndarray to PIL.Image format
Parameters:
image (np.ndarray):
Generated image
meta (Dict):
Metadata obtained on latents preparing step, can be empty
output_type (str, *optional*, pil):
Output format for result, can be pil or numpy
Returns:
image (List of np.ndarray or PIL.Image.Image):
Postprocessed images
if "src_height" in meta:
orig_height, orig_width = meta["src_height"], meta["src_width"]
image = [cv2.resize(img, (orig_width, orig_height))
for img in image]
return image
"""
if "padding" in meta:
pad = meta["padding"]
(_, end_h), (_, end_w) = pad[1:3]
h, w = image.shape[2:]
#print("image shape",image.shape[2:])
unpad_h = h - end_h
unpad_w = w - end_w
image = image[:, :, :unpad_h, :unpad_w]
image = np.clip(image / 2 + 0.5, 0, 1)
image = (image[0].transpose(1, 2, 0)[:, :, ::-1] * 255).astype(np.uint8)
if "src_height" in meta:
orig_height, orig_width = meta["src_height"], meta["src_width"]
image = cv2.resize(image, (orig_width, orig_height))
return image
def get_timesteps(self, num_inference_steps:int, strength:float, scheduler):
"""
Helper function for getting scheduler timesteps for generation
In case of image-to-image generation, it updates number of steps according to strength
Parameters:
num_inference_steps (int):
number of inference steps for generation
strength (float):
value between 0.0 and 1.0, that controls the amount of noise that is added to the input image.
Values that approach 1.0 allow for lots of variations but will also produce images that are not semantically consistent with the input.
"""
# get the original timestep using init_timestep
init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
t_start = max(num_inference_steps - init_timestep, 0)
timesteps = scheduler.timesteps[t_start:]
return timesteps, num_inference_steps - t_start
class StableDiffusionEngine(DiffusionPipeline):
def __init__(
self,
model="bes-dev/stable-diffusion-v1-4-openvino",
tokenizer="openai/clip-vit-large-patch14",
device=["CPU","CPU","CPU","CPU"]):
self.core = Core()
self.core.set_property({'CACHE_DIR': os.path.join(model, 'cache')})
self.batch_size = 2 if device[1] == device[2] and device[1] == "GPU" else 1
try_enable_npu_turbo(device, self.core)
try:
self.tokenizer = CLIPTokenizer.from_pretrained(model, local_files_only=True)
except Exception as e:
print("Local tokenizer not found. Attempting to download...")
self.tokenizer = self.download_tokenizer(tokenizer, model)
print("Loading models... ")
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
text_future = executor.submit(self.load_model, model, "text_encoder", device[0])
vae_de_future = executor.submit(self.load_model, model, "vae_decoder", device[3])
vae_en_future = executor.submit(self.load_model, model, "vae_encoder", device[3])
if self.batch_size == 1:
if "int8" not in model:
unet_future = executor.submit(self.load_model, model, "unet_bs1", device[1])
unet_neg_future = executor.submit(self.load_model, model, "unet_bs1", device[2]) if device[1] != device[2] else None
else:
unet_future = executor.submit(self.load_model, model, "unet_int8a16", device[1])
unet_neg_future = executor.submit(self.load_model, model, "unet_int8a16", device[2]) if device[1] != device[2] else None
else:
unet_future = executor.submit(self.load_model, model, "unet", device[1])
unet_neg_future = None
self.unet = unet_future.result()
self.unet_neg = unet_neg_future.result() if unet_neg_future else self.unet
self.text_encoder = text_future.result()
self.vae_decoder = vae_de_future.result()
self.vae_encoder = vae_en_future.result()
print("Text Device:", device[0])
print("unet Device:", device[1])
print("unet-neg Device:", device[2])
print("VAE Device:", device[3])
self._text_encoder_output = self.text_encoder.output(0)
self._unet_output = self.unet.output(0)
self._vae_d_output = self.vae_decoder.output(0)
self._vae_e_output = self.vae_encoder.output(0) if self.vae_encoder else None
self.unet_input_tensor_name = "sample" if 'sample' in self.unet.input(0).names else "latent_model_input"
if self.batch_size == 1:
self.infer_request = self.unet.create_infer_request()
self.infer_request_neg = self.unet_neg.create_infer_request()
self._unet_neg_output = self.unet_neg.output(0)
else:
self.infer_request = None
self.infer_request_neg = None
self._unet_neg_output = None
self.set_dimensions()
def load_model(self, model, model_name, device):
if "NPU" in device:
with open(os.path.join(model, f"{model_name}.blob"), "rb") as f:
return self.core.import_model(f.read(), device)
return self.core.compile_model(os.path.join(model, f"{model_name}.xml"), device)
def set_dimensions(self):
latent_shape = self.unet.input(self.unet_input_tensor_name).shape
if latent_shape[1] == 4:
self.height = latent_shape[2] * 8
self.width = latent_shape[3] * 8
else:
self.height = latent_shape[1] * 8
self.width = latent_shape[2] * 8
def __call__(
self,
prompt,
init_image=None,
negative_prompt=None,
scheduler=None,
strength=0.5,
num_inference_steps=32,
guidance_scale=7.5,
eta=0.0,
create_gif=False,
model=None,
callback=None,
callback_userdata=None
):
# extract condition
text_input = self.tokenizer(
prompt,
padding="max_length",
max_length=self.tokenizer.model_max_length,
truncation=True,
return_tensors="np",
)
text_embeddings = self.text_encoder(text_input.input_ids)[self._text_encoder_output]
# do classifier free guidance
do_classifier_free_guidance = guidance_scale > 1.0
if do_classifier_free_guidance:
if negative_prompt is None:
uncond_tokens = [""]
elif isinstance(negative_prompt, str):
uncond_tokens = [negative_prompt]
else:
uncond_tokens = negative_prompt
tokens_uncond = self.tokenizer(
uncond_tokens,
padding="max_length",
max_length=self.tokenizer.model_max_length, # truncation=True,
return_tensors="np"
)
uncond_embeddings = self.text_encoder(tokens_uncond.input_ids)[self._text_encoder_output]
text_embeddings = np.concatenate([uncond_embeddings, text_embeddings])
# set timesteps
accepts_offset = "offset" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
extra_set_kwargs = {}
if accepts_offset:
extra_set_kwargs["offset"] = 1
scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)
timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, scheduler)
latent_timestep = timesteps[:1]
# get the initial random noise unless the user supplied it
latents, meta = self.prepare_latents(init_image, latent_timestep, scheduler,model)
# 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(scheduler.step).parameters.keys())
extra_step_kwargs = {}
if accepts_eta:
extra_step_kwargs["eta"] = eta
if create_gif:
frames = []
for i, t in enumerate(self.progress_bar(timesteps)):
if callback:
callback(i, callback_userdata)
if self.batch_size == 1:
# expand the latents if we are doing classifier free guidance
noise_pred = []
latent_model_input = latents
#Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the Euler algorithm.
latent_model_input = scheduler.scale_model_input(latent_model_input, t)
latent_model_input_pos = latent_model_input
latent_model_input_neg = latent_model_input
if self.unet.input(self.unet_input_tensor_name).shape[1] != 4:
try:
latent_model_input_pos = latent_model_input_pos.permute(0,2,3,1)
except:
latent_model_input_pos = latent_model_input_pos.transpose(0,2,3,1)
if self.unet_neg.input(self.unet_input_tensor_name).shape[1] != 4:
try:
latent_model_input_neg = latent_model_input_neg.permute(0,2,3,1)
except:
latent_model_input_neg = latent_model_input_neg.transpose(0,2,3,1)
if "sample" in self.unet_input_tensor_name:
input_tens_neg_dict = {"sample" : latent_model_input_neg, "encoder_hidden_states": np.expand_dims(text_embeddings[0], axis=0), "timestep": np.expand_dims(np.float32(t), axis=0)}
input_tens_pos_dict = {"sample" : latent_model_input_pos, "encoder_hidden_states": np.expand_dims(text_embeddings[1], axis=0), "timestep": np.expand_dims(np.float32(t), axis=0)}
else:
input_tens_neg_dict = {"latent_model_input" : latent_model_input_neg, "encoder_hidden_states": np.expand_dims(text_embeddings[0], axis=0), "t": np.expand_dims(np.float32(t), axis=0)}
input_tens_pos_dict = {"latent_model_input" : latent_model_input_pos, "encoder_hidden_states": np.expand_dims(text_embeddings[1], axis=0), "t": np.expand_dims(np.float32(t), axis=0)}
self.infer_request_neg.start_async(input_tens_neg_dict)
self.infer_request.start_async(input_tens_pos_dict)
self.infer_request_neg.wait()
self.infer_request.wait()
noise_pred_neg = self.infer_request_neg.get_output_tensor(0)
noise_pred_pos = self.infer_request.get_output_tensor(0)
noise_pred.append(noise_pred_neg.data.astype(np.float32))
noise_pred.append(noise_pred_pos.data.astype(np.float32))
else:
latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents
latent_model_input = scheduler.scale_model_input(latent_model_input, t)
noise_pred = self.unet([latent_model_input, np.array(t, dtype=np.float32), text_embeddings])[self._unet_output]
if do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred[0], noise_pred[1]
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents = scheduler.step(torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs)["prev_sample"].numpy()
if create_gif:
frames.append(latents)
if callback:
callback(num_inference_steps, callback_userdata)
# scale and decode the image latents with vae
#if self.height == 512 and self.width == 512:
latents = 1 / 0.18215 * latents
image = self.vae_decoder(latents)[self._vae_d_output]
image = self.postprocess_image(image, meta)
return image
def prepare_latents(self, image: PIL.Image.Image = None, latent_timestep: torch.Tensor = None,
scheduler=LMSDiscreteScheduler,model=None):
"""
Function for getting initial latents for starting generation
Parameters:
image (PIL.Image.Image, *optional*, None):
Input image for generation, if not provided randon noise will be used as starting point
latent_timestep (torch.Tensor, *optional*, None):
Predicted by scheduler initial step for image generation, required for latent image mixing with nosie
Returns:
latents (np.ndarray):
Image encoded in latent space
"""
latents_shape = (1, 4, self.height // 8, self.width // 8)
noise = np.random.randn(*latents_shape).astype(np.float32)
if image is None:
#print("Image is NONE")
# if we use LMSDiscreteScheduler, let's make sure latents are mulitplied by sigmas
if isinstance(scheduler, LMSDiscreteScheduler):
noise = noise * scheduler.sigmas[0].numpy()
return noise, {}
elif isinstance(scheduler, EulerDiscreteScheduler):
noise = noise * scheduler.sigmas.max().numpy()
return noise, {}
else:
return noise, {}
input_image, meta = preprocess(image, self.height, self.width)
moments = self.vae_encoder(input_image)[self._vae_e_output]
if "sd_2.1" in model:
latents = moments * 0.18215
else:
mean, logvar = np.split(moments, 2, axis=1)
std = np.exp(logvar * 0.5)
latents = (mean + std * np.random.randn(*mean.shape)) * 0.18215
latents = scheduler.add_noise(torch.from_numpy(latents), torch.from_numpy(noise), latent_timestep).numpy()
return latents, meta
def postprocess_image(self, image: np.ndarray, meta: Dict):
"""
Postprocessing for decoded image. Takes generated image decoded by VAE decoder, unpad it to initila image size (if required),
normalize and convert to [0, 255] pixels range. Optionally, convertes it from np.ndarray to PIL.Image format
Parameters:
image (np.ndarray):
Generated image
meta (Dict):
Metadata obtained on latents preparing step, can be empty
output_type (str, *optional*, pil):
Output format for result, can be pil or numpy
Returns:
image (List of np.ndarray or PIL.Image.Image):
Postprocessed images
if "src_height" in meta:
orig_height, orig_width = meta["src_height"], meta["src_width"]
image = [cv2.resize(img, (orig_width, orig_height))
for img in image]
return image
"""
if "padding" in meta:
pad = meta["padding"]
(_, end_h), (_, end_w) = pad[1:3]
h, w = image.shape[2:]
# print("image shape",image.shape[2:])
unpad_h = h - end_h
unpad_w = w - end_w
image = image[:, :, :unpad_h, :unpad_w]
image = np.clip(image / 2 + 0.5, 0, 1)
image = (image[0].transpose(1, 2, 0)[:, :, ::-1] * 255).astype(np.uint8)
if "src_height" in meta:
orig_height, orig_width = meta["src_height"], meta["src_width"]
image = cv2.resize(image, (orig_width, orig_height))
return image
# image = (image / 2 + 0.5).clip(0, 1)
# image = (image[0].transpose(1, 2, 0)[:, :, ::-1] * 255).astype(np.uint8)
def get_timesteps(self, num_inference_steps: int, strength: float, scheduler):
"""
Helper function for getting scheduler timesteps for generation
In case of image-to-image generation, it updates number of steps according to strength
Parameters:
num_inference_steps (int):
number of inference steps for generation
strength (float):
value between 0.0 and 1.0, that controls the amount of noise that is added to the input image.
Values that approach 1.0 allow for lots of variations but will also produce images that are not semantically consistent with the input.
"""
# get the original timestep using init_timestep
init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
t_start = max(num_inference_steps - init_timestep, 0)
timesteps = scheduler.timesteps[t_start:]
return timesteps, num_inference_steps - t_start
class LatentConsistencyEngine(DiffusionPipeline):
def __init__(
self,
model="SimianLuo/LCM_Dreamshaper_v7",
tokenizer="openai/clip-vit-large-patch14",
device=["CPU", "CPU", "CPU"],
):
super().__init__()
try:
self.tokenizer = CLIPTokenizer.from_pretrained(model, local_files_only=True)
except:
self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer)
self.tokenizer.save_pretrained(model)
self.core = Core()
self.core.set_property({'CACHE_DIR': os.path.join(model, 'cache')}) # adding caching to reduce init time
try_enable_npu_turbo(device, self.core)
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
text_future = executor.submit(self.load_model, model, "text_encoder", device[0])
unet_future = executor.submit(self.load_model, model, "unet", device[1])
vae_de_future = executor.submit(self.load_model, model, "vae_decoder", device[2])
print("Text Device:", device[0])
self.text_encoder = text_future.result()
self._text_encoder_output = self.text_encoder.output(0)
print("Unet Device:", device[1])
self.unet = unet_future.result()
self._unet_output = self.unet.output(0)
self.infer_request = self.unet.create_infer_request()
print(f"VAE Device: {device[2]}")
self.vae_decoder = vae_de_future.result()
self.infer_request_vae = self.vae_decoder.create_infer_request()
self.safety_checker = None #pipe.safety_checker
self.feature_extractor = None #pipe.feature_extractor
self.vae_scale_factor = 2 ** 3
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
def load_model(self, model, model_name, device):
if "NPU" in device:
with open(os.path.join(model, f"{model_name}.blob"), "rb") as f:
return self.core.import_model(f.read(), device)
return self.core.compile_model(os.path.join(model, f"{model_name}.xml"), device)
def _encode_prompt(
self,
prompt,
num_images_per_prompt,
prompt_embeds: None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
prompt (`str` or `List[str]`, *optional*):
prompt to be encoded
num_images_per_prompt (`int`):
number of images that should be generated per prompt
prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
"""
if prompt_embeds is None:
text_inputs = self.tokenizer(
prompt,
padding="max_length",
max_length=self.tokenizer.model_max_length,
truncation=True,
return_tensors="pt",
)
text_input_ids = text_inputs.input_ids
untruncated_ids = self.tokenizer(
prompt, padding="longest", return_tensors="pt"
).input_ids
if untruncated_ids.shape[-1] >= text_input_ids.shape[
-1
] and not torch.equal(text_input_ids, untruncated_ids):
removed_text = self.tokenizer.batch_decode(
untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
)
logger.warning(
"The following part of your input was truncated because CLIP can only handle sequences up to"
f" {self.tokenizer.model_max_length} tokens: {removed_text}"
)
prompt_embeds = self.text_encoder(text_input_ids, share_inputs=True, share_outputs=True)
prompt_embeds = torch.from_numpy(prompt_embeds[0])
bs_embed, seq_len, _ = prompt_embeds.shape
# duplicate text embeddings for each generation per prompt
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
prompt_embeds = prompt_embeds.view(
bs_embed * num_images_per_prompt, seq_len, -1
)
# Don't need to get uncond prompt embedding because of LCM Guided Distillation
return prompt_embeds
def run_safety_checker(self, image, dtype):
if self.safety_checker is None:
has_nsfw_concept = None
else:
if torch.is_tensor(image):
feature_extractor_input = self.image_processor.postprocess(
image, output_type="pil"
)
else:
feature_extractor_input = self.image_processor.numpy_to_pil(image)
safety_checker_input = self.feature_extractor(
feature_extractor_input, return_tensors="pt"
)
image, has_nsfw_concept = self.safety_checker(
images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
)
return image, has_nsfw_concept
def prepare_latents(
self, batch_size, num_channels_latents, height, width, dtype, latents=None
):
shape = (
batch_size,
num_channels_latents,
height // self.vae_scale_factor,
width // self.vae_scale_factor,
)
if latents is None:
latents = torch.randn(shape, dtype=dtype)
# scale the initial noise by the standard deviation required by the scheduler
return latents
def get_w_embedding(self, w, embedding_dim=512, dtype=torch.float32):
"""
see https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
Args:
timesteps: torch.Tensor: generate embedding vectors at these timesteps
embedding_dim: int: dimension of the embeddings to generate
dtype: data type of the generated embeddings
Returns:
embedding vectors with shape `(len(timesteps), embedding_dim)`
"""
assert len(w.shape) == 1
w = w * 1000.0
half_dim = embedding_dim // 2
emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
emb = w.to(dtype)[:, None] * emb[None, :]
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
if embedding_dim % 2 == 1: # zero pad
emb = torch.nn.functional.pad(emb, (0, 1))
assert emb.shape == (w.shape[0], embedding_dim)
return emb
@torch.no_grad()
def __call__(
self,
prompt: Union[str, List[str]] = None,
height: Optional[int] = 512,
width: Optional[int] = 512,
guidance_scale: float = 7.5,
scheduler = None,
num_images_per_prompt: Optional[int] = 1,
latents: Optional[torch.FloatTensor] = None,
num_inference_steps: int = 4,
lcm_origin_steps: int = 50,
prompt_embeds: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
model: Optional[Dict[str, any]] = None,
seed: Optional[int] = 1234567,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
callback = None,
callback_userdata = None
):
# 1. 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 seed is not None:
torch.manual_seed(seed)
#print("After Step 1: batch size is ", batch_size)
# do_classifier_free_guidance = guidance_scale > 0.0
# In LCM Implementation: cfg_noise = noise_cond + cfg_scale * (noise_cond - noise_uncond) , (cfg_scale > 0.0 using CFG)
# 2. Encode input prompt
prompt_embeds = self._encode_prompt(
prompt,
num_images_per_prompt,
prompt_embeds=prompt_embeds,
)
#print("After Step 2: prompt embeds is ", prompt_embeds)
#print("After Step 2: scheduler is ", scheduler )
# 3. Prepare timesteps
scheduler.set_timesteps(num_inference_steps, original_inference_steps=lcm_origin_steps)
timesteps = scheduler.timesteps
#print("After Step 3: timesteps is ", timesteps)
# 4. Prepare latent variable
num_channels_latents = 4
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
prompt_embeds.dtype,
latents,
)
latents = latents * scheduler.init_noise_sigma
#print("After Step 4: ")
bs = batch_size * num_images_per_prompt
# 5. Get Guidance Scale Embedding
w = torch.tensor(guidance_scale).repeat(bs)
w_embedding = self.get_w_embedding(w, embedding_dim=256)
#print("After Step 5: ")
# 6. LCM MultiStep Sampling Loop:
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
if callback:
callback(i+1, callback_userdata)
ts = torch.full((bs,), t, dtype=torch.long)
# model prediction (v-prediction, eps, x)
model_pred = self.unet([latents, ts, prompt_embeds, w_embedding],share_inputs=True, share_outputs=True)[0]
# compute the previous noisy sample x_t -> x_t-1
latents, denoised = scheduler.step(
torch.from_numpy(model_pred), t, latents, return_dict=False
)
progress_bar.update()
#print("After Step 6: ")
vae_start = time.time()
if not output_type == "latent":
image = torch.from_numpy(self.vae_decoder(denoised / 0.18215, share_inputs=True, share_outputs=True)[0])
else:
image = denoised
print("Decoder Ended: ", time.time() - vae_start)
#post_start = time.time()
#if has_nsfw_concept is None:
do_denormalize = [True] * image.shape[0]
#else:
# do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
#print ("After do_denormalize: image is ", image)
image = self.image_processor.postprocess(
image, output_type=output_type, do_denormalize=do_denormalize
)
return image[0]
class LatentConsistencyEngineAdvanced(DiffusionPipeline):
def __init__(
self,
model="SimianLuo/LCM_Dreamshaper_v7",
tokenizer="openai/clip-vit-large-patch14",
device=["CPU", "CPU", "CPU"],
):
super().__init__()
try:
self.tokenizer = CLIPTokenizer.from_pretrained(model, local_files_only=True)
except:
self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer)
self.tokenizer.save_pretrained(model)
self.core = Core()
self.core.set_property({'CACHE_DIR': os.path.join(model, 'cache')}) # adding caching to reduce init time
#try_enable_npu_turbo(device, self.core)
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
text_future = executor.submit(self.load_model, model, "text_encoder", device[0])
unet_future = executor.submit(self.load_model, model, "unet", device[1])
vae_de_future = executor.submit(self.load_model, model, "vae_decoder", device[2])
vae_encoder_future = executor.submit(self.load_model, model, "vae_encoder", device[2])
print("Text Device:", device[0])
self.text_encoder = text_future.result()
self._text_encoder_output = self.text_encoder.output(0)
print("Unet Device:", device[1])
self.unet = unet_future.result()
self._unet_output = self.unet.output(0)
self.infer_request = self.unet.create_infer_request()
print(f"VAE Device: {device[2]}")
self.vae_decoder = vae_de_future.result()
self.vae_encoder = vae_encoder_future.result()
self._vae_e_output = self.vae_encoder.output(0) if self.vae_encoder else None
self.infer_request_vae = self.vae_decoder.create_infer_request()
self.safety_checker = None #pipe.safety_checker
self.feature_extractor = None #pipe.feature_extractor
self.vae_scale_factor = 2 ** 3
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
def load_model(self, model, model_name, device):
print(f"Compiling the {model_name} to {device} ...")
return self.core.compile_model(os.path.join(model, f"{model_name}.xml"), device)
def get_timesteps(self, num_inference_steps:int, strength:float, scheduler):
"""
Helper function for getting scheduler timesteps for generation
In case of image-to-image generation, it updates number of steps according to strength
Parameters:
num_inference_steps (int):
number of inference steps for generation
strength (float):
value between 0.0 and 1.0, that controls the amount of noise that is added to the input image.
Values that approach 1.0 allow for lots of variations but will also produce images that are not semantically consistent with the input.
"""
# get the original timestep using init_timestep
init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
t_start = max(num_inference_steps - init_timestep, 0)
timesteps = scheduler.timesteps[t_start:]
return timesteps, num_inference_steps - t_start
def _encode_prompt(
self,
prompt,
num_images_per_prompt,
prompt_embeds: None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
prompt (`str` or `List[str]`, *optional*):
prompt to be encoded
num_images_per_prompt (`int`):
number of images that should be generated per prompt
prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
"""
if prompt_embeds is None:
text_inputs = self.tokenizer(
prompt,
padding="max_length",
max_length=self.tokenizer.model_max_length,
truncation=True,
return_tensors="pt",
)
text_input_ids = text_inputs.input_ids
untruncated_ids = self.tokenizer(
prompt, padding="longest", return_tensors="pt"
).input_ids
if untruncated_ids.shape[-1] >= text_input_ids.shape[
-1
] and not torch.equal(text_input_ids, untruncated_ids):
removed_text = self.tokenizer.batch_decode(
untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
)
logger.warning(
"The following part of your input was truncated because CLIP can only handle sequences up to"
f" {self.tokenizer.model_max_length} tokens: {removed_text}"
)
prompt_embeds = self.text_encoder(text_input_ids, share_inputs=True, share_outputs=True)
prompt_embeds = torch.from_numpy(prompt_embeds[0])
bs_embed, seq_len, _ = prompt_embeds.shape
# duplicate text embeddings for each generation per prompt
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
prompt_embeds = prompt_embeds.view(
bs_embed * num_images_per_prompt, seq_len, -1
)
# Don't need to get uncond prompt embedding because of LCM Guided Distillation
return prompt_embeds
def run_safety_checker(self, image, dtype):
if self.safety_checker is None:
has_nsfw_concept = None
else:
if torch.is_tensor(image):
feature_extractor_input = self.image_processor.postprocess(
image, output_type="pil"
)
else:
feature_extractor_input = self.image_processor.numpy_to_pil(image)
safety_checker_input = self.feature_extractor(
feature_extractor_input, return_tensors="pt"
)
image, has_nsfw_concept = self.safety_checker(
images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
)
return image, has_nsfw_concep
def prepare_latents(
self,image,timestep,batch_size, num_channels_latents, height, width, dtype, scheduler,latents=None,
):
shape = (
batch_size,
num_channels_latents,
height // self.vae_scale_factor,
width // self.vae_scale_factor,
)
if image:
#latents_shape = (1, 4, 512, 512 // 8)
#input_image, meta = preprocess(image,512,512)
latents_shape = (1, 4, 512 // 8, 512 // 8)
noise = np.random.randn(*latents_shape).astype(np.float32)
input_image,meta = preprocess(image,512,512)
moments = self.vae_encoder(input_image)[self._vae_e_output]
mean, logvar = np.split(moments, 2, axis=1)
std = np.exp(logvar * 0.5)
latents = (mean + std * np.random.randn(*mean.shape)) * 0.18215
noise = torch.randn(shape, dtype=dtype)
#latents = scheduler.add_noise(init_latents, noise, timestep)
latents = scheduler.add_noise(torch.from_numpy(latents), noise, timestep)
else:
latents = torch.randn(shape, dtype=dtype)
# scale the initial noise by the standard deviation required by the scheduler
return latents
def get_w_embedding(self, w, embedding_dim=512, dtype=torch.float32):
"""
see https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
Args:
timesteps: torch.Tensor: generate embedding vectors at these timesteps
embedding_dim: int: dimension of the embeddings to generate
dtype: data type of the generated embeddings
Returns:
embedding vectors with shape `(len(timesteps), embedding_dim)`
"""
assert len(w.shape) == 1
w = w * 1000.0
half_dim = embedding_dim // 2
emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
emb = w.to(dtype)[:, None] * emb[None, :]
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
if embedding_dim % 2 == 1: # zero pad
emb = torch.nn.functional.pad(emb, (0, 1))
assert emb.shape == (w.shape[0], embedding_dim)
return emb
@torch.no_grad()
def __call__(
self,
prompt: Union[str, List[str]] = None,
init_image: Optional[PIL.Image.Image] = None,
strength: Optional[float] = 0.8,
height: Optional[int] = 512,
width: Optional[int] = 512,
guidance_scale: float = 7.5,
scheduler = None,
num_images_per_prompt: Optional[int] = 1,
latents: Optional[torch.FloatTensor] = None,
num_inference_steps: int = 4,
lcm_origin_steps: int = 50,
prompt_embeds: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
model: Optional[Dict[str, any]] = None,
seed: Optional[int] = 1234567,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
callback = None,
callback_userdata = None
):
# 1. 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 seed is not None:
torch.manual_seed(seed)
#print("After Step 1: batch size is ", batch_size)
# do_classifier_free_guidance = guidance_scale > 0.0
# In LCM Implementation: cfg_noise = noise_cond + cfg_scale * (noise_cond - noise_uncond) , (cfg_scale > 0.0 using CFG)
# 2. Encode input prompt
prompt_embeds = self._encode_prompt(
prompt,
num_images_per_prompt,
prompt_embeds=prompt_embeds,
)
#print("After Step 2: prompt embeds is ", prompt_embeds)
#print("After Step 2: scheduler is ", scheduler )
# 3. Prepare timesteps
#scheduler.set_timesteps(num_inference_steps, original_inference_steps=lcm_origin_steps)
latent_timestep = None
if init_image:
scheduler.set_timesteps(num_inference_steps, original_inference_steps=lcm_origin_steps)
timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, scheduler)
latent_timestep = timesteps[:1]
else:
scheduler.set_timesteps(num_inference_steps, original_inference_steps=lcm_origin_steps)
timesteps = scheduler.timesteps
#timesteps = scheduler.timesteps
#latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
#print("timesteps: ", latent_timestep)
#print("After Step 3: timesteps is ", timesteps)
# 4. Prepare latent variable
num_channels_latents = 4
latents = self.prepare_latents(
init_image,
latent_timestep,
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
prompt_embeds.dtype,
scheduler,
latents,
)
latents = latents * scheduler.init_noise_sigma
#print("After Step 4: ")
bs = batch_size * num_images_per_prompt
# 5. Get Guidance Scale Embedding
w = torch.tensor(guidance_scale).repeat(bs)
w_embedding = self.get_w_embedding(w, embedding_dim=256)
#print("After Step 5: ")
# 6. LCM MultiStep Sampling Loop:
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
if callback:
callback(i+1, callback_userdata)
ts = torch.full((bs,), t, dtype=torch.long)
# model prediction (v-prediction, eps, x)
model_pred = self.unet([latents, ts, prompt_embeds, w_embedding],share_inputs=True, share_outputs=True)[0]
# compute the previous noisy sample x_t -> x_t-1
latents, denoised = scheduler.step(
torch.from_numpy(model_pred), t, latents, return_dict=False
)
progress_bar.update()
#print("After Step 6: ")
vae_start = time.time()
if not output_type == "latent":
image = torch.from_numpy(self.vae_decoder(denoised / 0.18215, share_inputs=True, share_outputs=True)[0])
else:
image = denoised
print("Decoder Ended: ", time.time() - vae_start)
#post_start = time.time()
#if has_nsfw_concept is None:
do_denormalize = [True] * image.shape[0]
#else:
# do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
#print ("After do_denormalize: image is ", image)
image = self.image_processor.postprocess(
image, output_type=output_type, do_denormalize=do_denormalize
)
return image[0]
class StableDiffusionEngineReferenceOnly(DiffusionPipeline):
def __init__(
self,
#scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
model="bes-dev/stable-diffusion-v1-4-openvino",
tokenizer="openai/clip-vit-large-patch14",
device=["CPU","CPU","CPU"]
):
#self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer)
try:
self.tokenizer = CLIPTokenizer.from_pretrained(model,local_files_only=True)
except:
self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer)
self.tokenizer.save_pretrained(model)
#self.scheduler = scheduler
# models
self.core = Core()
self.core.set_property({'CACHE_DIR': os.path.join(model, 'cache')}) #adding caching to reduce init time
# text features
print("Text Device:",device[0])
self.text_encoder = self.core.compile_model(os.path.join(model, "text_encoder.xml"), device[0])
self._text_encoder_output = self.text_encoder.output(0)
# diffusion
print("unet_w Device:",device[1])
self.unet_w = self.core.compile_model(os.path.join(model, "unet_reference_write.xml"), device[1])
self._unet_w_output = self.unet_w.output(0)
self.latent_shape = tuple(self.unet_w.inputs[0].shape)[1:]
print("unet_r Device:",device[1])
self.unet_r = self.core.compile_model(os.path.join(model, "unet_reference_read.xml"), device[1])
self._unet_r_output = self.unet_r.output(0)
# decoder
print("Vae Device:",device[2])
self.vae_decoder = self.core.compile_model(os.path.join(model, "vae_decoder.xml"), device[2])
# encoder
self.vae_encoder = self.core.compile_model(os.path.join(model, "vae_encoder.xml"), device[2])
self.init_image_shape = tuple(self.vae_encoder.inputs[0].shape)[2:]
self._vae_d_output = self.vae_decoder.output(0)
self._vae_e_output = self.vae_encoder.output(0) if self.vae_encoder is not None else None
self.height = self.unet_w.input(0).shape[2] * 8
self.width = self.unet_w.input(0).shape[3] * 8
def __call__(
self,
prompt,
image = None,
negative_prompt=None,
scheduler=None,
strength = 1.0,
num_inference_steps = 32,
guidance_scale = 7.5,
eta = 0.0,
create_gif = False,
model = None,
callback = None,
callback_userdata = None
):
# extract condition
text_input = self.tokenizer(
prompt,
padding="max_length",
max_length=self.tokenizer.model_max_length,
truncation=True,
return_tensors="np",
)
text_embeddings = self.text_encoder(text_input.input_ids)[self._text_encoder_output]
# do classifier free guidance
do_classifier_free_guidance = guidance_scale > 1.0
if do_classifier_free_guidance:
if negative_prompt is None:
uncond_tokens = [""]
elif isinstance(negative_prompt, str):
uncond_tokens = [negative_prompt]
else:
uncond_tokens = negative_prompt
tokens_uncond = self.tokenizer(
uncond_tokens,
padding="max_length",
max_length=self.tokenizer.model_max_length, #truncation=True,
return_tensors="np"
)
uncond_embeddings = self.text_encoder(tokens_uncond.input_ids)[self._text_encoder_output]
text_embeddings = np.concatenate([uncond_embeddings, text_embeddings])
# set timesteps
accepts_offset = "offset" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
extra_set_kwargs = {}
if accepts_offset:
extra_set_kwargs["offset"] = 1
scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)
timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, scheduler)
latent_timestep = timesteps[:1]
ref_image = self.prepare_image(
image=image,
width=512,
height=512,
)
# get the initial random noise unless the user supplied it
latents, meta = self.prepare_latents(None, latent_timestep, scheduler)
#ref_image_latents, _ = self.prepare_latents(init_image, latent_timestep, scheduler)
ref_image_latents = self.ov_prepare_ref_latents(ref_image)
# 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(scheduler.step).parameters.keys())
extra_step_kwargs = {}
if accepts_eta:
extra_step_kwargs["eta"] = eta
if create_gif:
frames = []
for i, t in enumerate(self.progress_bar(timesteps)):
if callback:
callback(i, callback_userdata)
# 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 = scheduler.scale_model_input(latent_model_input, t)
# ref only part
noise = randn_tensor(
ref_image_latents.shape
)
ref_xt = scheduler.add_noise(
torch.from_numpy(ref_image_latents),
noise,
t.reshape(
1,
),
).numpy()
ref_xt = np.concatenate([ref_xt] * 2) if do_classifier_free_guidance else ref_xt
ref_xt = scheduler.scale_model_input(ref_xt, t)
# MODE = "write"
result_w_dict = self.unet_w([
ref_xt,
t,
text_embeddings
])
down_0_attn0 = result_w_dict["/unet/down_blocks.0/attentions.0/transformer_blocks.0/norm1/LayerNormalization_output_0"]
down_0_attn1 = result_w_dict["/unet/down_blocks.0/attentions.1/transformer_blocks.0/norm1/LayerNormalization_output_0"]
down_1_attn0 = result_w_dict["/unet/down_blocks.1/attentions.0/transformer_blocks.0/norm1/LayerNormalization_output_0"]
down_1_attn1 = result_w_dict["/unet/down_blocks.1/attentions.1/transformer_blocks.0/norm1/LayerNormalization_output_0"]
down_2_attn0 = result_w_dict["/unet/down_blocks.2/attentions.0/transformer_blocks.0/norm1/LayerNormalization_output_0"]
down_2_attn1 = result_w_dict["/unet/down_blocks.2/attentions.1/transformer_blocks.0/norm1/LayerNormalization_output_0"]
mid_attn0 = result_w_dict["/unet/mid_block/attentions.0/transformer_blocks.0/norm1/LayerNormalization_output_0"]
up_1_attn0 = result_w_dict["/unet/up_blocks.1/attentions.0/transformer_blocks.0/norm1/LayerNormalization_output_0"]
up_1_attn1 = result_w_dict["/unet/up_blocks.1/attentions.1/transformer_blocks.0/norm1/LayerNormalization_output_0"]
up_1_attn2 = result_w_dict["/unet/up_blocks.1/attentions.2/transformer_blocks.0/norm1/LayerNormalization_output_0"]
up_2_attn0 = result_w_dict["/unet/up_blocks.2/attentions.0/transformer_blocks.0/norm1/LayerNormalization_output_0"]
up_2_attn1 = result_w_dict["/unet/up_blocks.2/attentions.1/transformer_blocks.0/norm1/LayerNormalization_output_0"]
up_2_attn2 = result_w_dict["/unet/up_blocks.2/attentions.2/transformer_blocks.0/norm1/LayerNormalization_output_0"]
up_3_attn0 = result_w_dict["/unet/up_blocks.3/attentions.0/transformer_blocks.0/norm1/LayerNormalization_output_0"]
up_3_attn1 = result_w_dict["/unet/up_blocks.3/attentions.1/transformer_blocks.0/norm1/LayerNormalization_output_0"]
up_3_attn2 = result_w_dict["/unet/up_blocks.3/attentions.2/transformer_blocks.0/norm1/LayerNormalization_output_0"]
# MODE = "read"
noise_pred = self.unet_r([
latent_model_input, t, text_embeddings, down_0_attn0, down_0_attn1, down_1_attn0,
down_1_attn1, down_2_attn0, down_2_attn1, mid_attn0, up_1_attn0, up_1_attn1, up_1_attn2,
up_2_attn0, up_2_attn1, up_2_attn2, up_3_attn0, up_3_attn1, up_3_attn2
])[0]
# perform guidance
if do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred[0], noise_pred[1]
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents = scheduler.step(torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs)["prev_sample"].numpy()
if create_gif:
frames.append(latents)
if callback:
callback(num_inference_steps, callback_userdata)
# scale and decode the image latents with vae
image = self.vae_decoder(latents)[self._vae_d_output]
image = self.postprocess_image(image, meta)
if create_gif:
gif_folder=os.path.join(model,"../../../gif")
if not os.path.exists(gif_folder):
os.makedirs(gif_folder)
for i in range(0,len(frames)):
image = self.vae_decoder(frames[i])[self._vae_d_output]
image = self.postprocess_image(image, meta)
output = gif_folder + "/" + str(i).zfill(3) +".png"
cv2.imwrite(output, image)
with open(os.path.join(gif_folder, "prompt.json"), "w") as file:
json.dump({"prompt": prompt}, file)
frames_image = [Image.open(image) for image in glob.glob(f"{gif_folder}/*.png")]
frame_one = frames_image[0]
gif_file=os.path.join(gif_folder,"stable_diffusion.gif")
frame_one.save(gif_file, format="GIF", append_images=frames_image, save_all=True, duration=100, loop=0)
return image
def ov_prepare_ref_latents(self, refimage, vae_scaling_factor=0.18215):
#refimage = refimage.to(device=device, dtype=dtype)
# encode the mask image into latents space so we can concatenate it to the latents
moments = self.vae_encoder(refimage)[0]
mean, logvar = np.split(moments, 2, axis=1)
std = np.exp(logvar * 0.5)
ref_image_latents = (mean + std * np.random.randn(*mean.shape))
ref_image_latents = vae_scaling_factor * ref_image_latents
#ref_image_latents = scheduler.add_noise(torch.from_numpy(ref_image_latents), torch.from_numpy(noise), latent_timestep).numpy()
# aligning device to prevent device errors when concating it with the latent model input
#ref_image_latents = ref_image_latents.to(device=device, dtype=dtype)
return ref_image_latents
def prepare_latents(self, image:PIL.Image.Image = None, latent_timestep:torch.Tensor = None, scheduler = LMSDiscreteScheduler):
"""
Function for getting initial latents for starting generation
Parameters:
image (PIL.Image.Image, *optional*, None):
Input image for generation, if not provided randon noise will be used as starting point
latent_timestep (torch.Tensor, *optional*, None):
Predicted by scheduler initial step for image generation, required for latent image mixing with nosie
Returns:
latents (np.ndarray):
Image encoded in latent space
"""
latents_shape = (1, 4, self.height // 8, self.width // 8)
noise = np.random.randn(*latents_shape).astype(np.float32)
if image is None:
#print("Image is NONE")
# if we use LMSDiscreteScheduler, let's make sure latents are mulitplied by sigmas
if isinstance(scheduler, LMSDiscreteScheduler):
noise = noise * scheduler.sigmas[0].numpy()
return noise, {}
elif isinstance(scheduler, EulerDiscreteScheduler):
noise = noise * scheduler.sigmas.max().numpy()
return noise, {}
else:
return noise, {}
input_image, meta = preprocess(image,self.height,self.width)
moments = self.vae_encoder(input_image)[self._vae_e_output]
mean, logvar = np.split(moments, 2, axis=1)
std = np.exp(logvar * 0.5)
latents = (mean + std * np.random.randn(*mean.shape)) * 0.18215
latents = scheduler.add_noise(torch.from_numpy(latents), torch.from_numpy(noise), latent_timestep).numpy()
return latents, meta
def postprocess_image(self, image:np.ndarray, meta:Dict):
"""
Postprocessing for decoded image. Takes generated image decoded by VAE decoder, unpad it to initila image size (if required),
normalize and convert to [0, 255] pixels range. Optionally, convertes it from np.ndarray to PIL.Image format
Parameters:
image (np.ndarray):
Generated image
meta (Dict):
Metadata obtained on latents preparing step, can be empty
output_type (str, *optional*, pil):
Output format for result, can be pil or numpy
Returns:
image (List of np.ndarray or PIL.Image.Image):
Postprocessed images
if "src_height" in meta:
orig_height, orig_width = meta["src_height"], meta["src_width"]
image = [cv2.resize(img, (orig_width, orig_height))
for img in image]
return image
"""
if "padding" in meta:
pad = meta["padding"]
(_, end_h), (_, end_w) = pad[1:3]
h, w = image.shape[2:]
#print("image shape",image.shape[2:])
unpad_h = h - end_h
unpad_w = w - end_w
image = image[:, :, :unpad_h, :unpad_w]
image = np.clip(image / 2 + 0.5, 0, 1)
image = (image[0].transpose(1, 2, 0)[:, :, ::-1] * 255).astype(np.uint8)
if "src_height" in meta:
orig_height, orig_width = meta["src_height"], meta["src_width"]
image = cv2.resize(image, (orig_width, orig_height))
return image
#image = (image / 2 + 0.5).clip(0, 1)
#image = (image[0].transpose(1, 2, 0)[:, :, ::-1] * 255).astype(np.uint8)
def get_timesteps(self, num_inference_steps:int, strength:float, scheduler):
"""
Helper function for getting scheduler timesteps for generation
In case of image-to-image generation, it updates number of steps according to strength
Parameters:
num_inference_steps (int):
number of inference steps for generation
strength (float):
value between 0.0 and 1.0, that controls the amount of noise that is added to the input image.
Values that approach 1.0 allow for lots of variations but will also produce images that are not semantically consistent with the input.
"""
# get the original timestep using init_timestep
init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
t_start = max(num_inference_steps - init_timestep, 0)
timesteps = scheduler.timesteps[t_start:]
return timesteps, num_inference_steps - t_start
def prepare_image(
self,
image,
width,
height,
do_classifier_free_guidance=False,
guess_mode=False,
):
if not isinstance(image, np.ndarray):
if isinstance(image, PIL.Image.Image):
image = [image]
if isinstance(image[0], PIL.Image.Image):
images = []
for image_ in image:
image_ = image_.convert("RGB")
image_ = image_.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])
image_ = np.array(image_)
image_ = image_[None, :]
images.append(image_)
image = images
image = np.concatenate(image, axis=0)
image = np.array(image).astype(np.float32) / 255.0
image = (image - 0.5) / 0.5
image = image.transpose(0, 3, 1, 2)
elif isinstance(image[0], np.ndarray):
image = np.concatenate(image, dim=0)
if do_classifier_free_guidance and not guess_mode:
image = np.concatenate([image] * 2)
return image
def print_npu_turbo_art():
random_number = random.randint(1, 3)
if random_number == 1:
print(" ")
print(" ___ ___ ___ ___ ___ ___ ")
print(" /\ \ /\ \ /\ \ /\ \ /\ \ _____ /\ \ ")
print(" \:\ \ /::\ \ \:\ \ ___ \:\ \ /::\ \ /::\ \ /::\ \ ")
print(" \:\ \ /:/\:\__\ \:\ \ /\__\ \:\ \ /:/\:\__\ /:/\:\ \ /:/\:\ \ ")
print(" _____\:\ \ /:/ /:/ / ___ \:\ \ /:/ / ___ \:\ \ /:/ /:/ / /:/ /::\__\ /:/ \:\ \ ")
print(" /::::::::\__\ /:/_/:/ / /\ \ \:\__\ /:/__/ /\ \ \:\__\ /:/_/:/__/___ /:/_/:/\:|__| /:/__/ \:\__\ ")
print(" \:\~~\~~\/__/ \:\/:/ / \:\ \ /:/ / /::\ \ \:\ \ /:/ / \:\/:::::/ / \:\/:/ /:/ / \:\ \ /:/ / ")
print(" \:\ \ \::/__/ \:\ /:/ / /:/\:\ \ \:\ /:/ / \::/~~/~~~~ \::/_/:/ / \:\ /:/ / ")
print(" \:\ \ \:\ \ \:\/:/ / \/__\:\ \ \:\/:/ / \:\~~\ \:\/:/ / \:\/:/ / ")
print(" \:\__\ \:\__\ \::/ / \:\__\ \::/ / \:\__\ \::/ / \::/ / ")
print(" \/__/ \/__/ \/__/ \/__/ \/__/ \/__/ \/__/ \/__/ ")
print(" ")
elif random_number == 2:
print(" _ _ ____ _ _ _____ _ _ ____ ____ ___ ")
print("| \ | | | _ \ | | | | |_ _| | | | | | _ \ | __ ) / _ \ ")
print("| \| | | |_) | | | | | | | | | | | | |_) | | _ \ | | | |")
print("| |\ | | __/ | |_| | | | | |_| | | _ < | |_) | | |_| |")
print("|_| \_| |_| \___/ |_| \___/ |_| \_\ |____/ \___/ ")
print(" ")
else:
print("")
print(" ) ( ( ) ")
print(" ( /( )\ ) * ) )\ ) ( ( /( ")
print(" )\()) (()/( ( ` ) /( ( (()/( ( )\ )\()) ")
print("((_)\ /(_)) )\ ( )(_)) )\ /(_)) )((_) ((_)\ ")
print(" _((_) (_)) _ ((_) (_(_()) _ ((_) (_)) ((_)_ ((_) ")
print("| \| | | _ \ | | | | |_ _| | | | | | _ \ | _ ) / _ \ ")
print("| .` | | _/ | |_| | | | | |_| | | / | _ \ | (_) | ")
print("|_|\_| |_| \___/ |_| \___/ |_|_\ |___/ \___/ ")
print(" ")