added files

app.py

@@ -1,7 +1,125 @@
-import gradio as gr
+from app_settings import AppSettings
+from utils import show_system_info
+import constants
+from argparse import ArgumentParser
+from context import Context
+from constants import APP_VERSION, LCM_DEFAULT_MODEL_OPENVINO
+from models.interface_types import InterfaceType
+from constants import DEVICE
 
-def greet(name):
-    return "Hello " + name + "!!"
+parser = ArgumentParser(description=f"FAST SD CPU {constants.APP_VERSION}")
+parser.add_argument(
+    "-s",
+    "--share",
+    action="store_true",
+    help="Create sharable link(Web UI)",
+    required=False,
+)
+group = parser.add_mutually_exclusive_group(required=False)
+group.add_argument(
+    "-g",
+    "--gui",
+    action="store_true",
+    help="Start desktop GUI",
+)
+group.add_argument(
+    "-w",
+    "--webui",
+    action="store_true",
+    help="Start Web UI",
+)
+group.add_argument(
+    "-v",
+    "--version",
+    action="store_true",
+    help="Version",
+)
+parser.add_argument(
+    "--lcm_model_id",
+    type=str,
+    help="Model ID or path,Default SimianLuo/LCM_Dreamshaper_v7",
+    default="SimianLuo/LCM_Dreamshaper_v7",
+)
+parser.add_argument(
+    "--prompt",
+    type=str,
+    help="Describe the image you want to generate",
+)
+parser.add_argument(
+    "--image_height",
+    type=int,
+    help="Height of the image",
+    default=512,
+)
+parser.add_argument(
+    "--image_width",
+    type=int,
+    help="Width of the image",
+    default=512,
+)
+parser.add_argument(
+    "--inference_steps",
+    type=int,
+    help="Number of steps,default : 4",
+    default=4,
+)
+parser.add_argument(
+    "--guidance_scale",
+    type=int,
+    help="Guidance scale,default : 8.0",
+    default=8.0,
+)
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+parser.add_argument(
+    "--number_of_images",
+    type=int,
+    help="Number of images to generate ,default : 1",
+    default=1,
+)
+parser.add_argument(
+    "--seed",
+    type=int,
+    help="Seed,default : -1 (disabled) ",
+    default=-1,
+)
+parser.add_argument(
+    "--use_openvino",
+    action="store_true",
+    help="Use OpenVINO model",
+)
+
+parser.add_argument(
+    "--use_offline_model",
+    action="store_true",
+    help="Use offline model",
+)
+parser.add_argument(
+    "--use_safety_checker",
+    action="store_false",
+    help="Use safety checker",
+)
+
+parser.add_argument(
+    "-i",
+    "--interactive",
+    action="store_true",
+    help="Interactive CLI mode",
+)
+
+args = parser.parse_args()
+
+if args.version:
+    print(APP_VERSION)
+    exit()
+
+parser.print_help()
+show_system_info()
+print(f"Using device : {constants.DEVICE}")
+
+app_settings = AppSettings()
+app_settings.load()
+
+
+start_webui(
+        app_settings,
+        args.share,)

app_settings.py

@@ -0,0 +1,46 @@
+import yaml
+from os import path, makedirs
+from models.settings import Settings
+from paths import FastStableDiffusionPaths
+
+
+class AppSettings:
+    def __init__(self):
+        self.config_path = FastStableDiffusionPaths().get_app_settings_path()
+
+    @property
+    def settings(self):
+        return self._config
+
+    def load(self):
+        if not path.exists(self.config_path):
+            base_dir = path.dirname(self.config_path)
+            if not path.exists(base_dir):
+                makedirs(base_dir)
+            try:
+                print("Settings not found creating default settings")
+                with open(self.config_path, "w") as file:
+                    yaml.dump(
+                        self._load_default(),
+                        file,
+                    )
+            except Exception as ex:
+                print(f"Error in creating settings : {ex}")
+                exit()
+        try:
+            with open(self.config_path) as file:
+                settings_dict = yaml.safe_load(file)
+                self._config = Settings.parse_obj(settings_dict)
+        except Exception as ex:
+            print(f"Error in loading settings : {ex}")
+
+    def save(self):
+        try:
+            with open(self.config_path, "w") as file:
+                yaml.dump(self._config.dict(), file)
+        except Exception as ex:
+            print(f"Error in saving settings : {ex}")
+
+    def _load_default(self) -> dict:
+        defult_config = Settings()
+        return defult_config.dict()

backend/image_saver.py

@@ -0,0 +1,39 @@
+from os import path, mkdir
+from typing import Any
+from uuid import uuid4
+from backend.models.lcmdiffusion_setting import LCMDiffusionSetting
+import json
+
+
+class ImageSaver:
+    @staticmethod
+    def save_images(
+        output_path: str,
+        images: Any,
+        folder_name: str = "",
+        format: str = ".png",
+        lcm_diffusion_setting: LCMDiffusionSetting = None,
+    ) -> None:
+        gen_id = uuid4()
+        for index, image in enumerate(images):
+            if not path.exists(output_path):
+                mkdir(output_path)
+
+            if folder_name:
+                out_path = path.join(
+                    output_path,
+                    folder_name,
+                )
+            else:
+                out_path = output_path
+
+            if not path.exists(out_path):
+                mkdir(out_path)
+            image.save(path.join(out_path, f"{gen_id}-{index+1}{format}"))
+        if lcm_diffusion_setting:
+            with open(path.join(out_path, f"{gen_id}.json"), "w") as json_file:
+                json.dump(
+                    lcm_diffusion_setting.model_dump(),
+                    json_file,
+                    indent=4,
+                )

backend/lcm_text_to_image.py

@@ -0,0 +1,115 @@
+from typing import Any
+from diffusers import DiffusionPipeline
+from os import path
+import torch
+from backend.models.lcmdiffusion_setting import LCMDiffusionSetting
+import numpy as np
+from constants import DEVICE
+
+
+if DEVICE == "cpu":
+    from backend.lcmdiffusion.pipelines.openvino.lcm_ov_pipeline import (
+        OVLatentConsistencyModelPipeline,
+    )
+    from backend.lcmdiffusion.pipelines.openvino.lcm_scheduler import (
+        LCMScheduler,
+    )
+
+
+class LCMTextToImage:
+    def __init__(
+        self,
+        device: str = "cpu",
+    ) -> None:
+        self.pipeline = None
+        self.use_openvino = False
+        self.device = None
+        self.previous_model_id = None
+
+    def _get_lcm_diffusion_pipeline_path(self) -> str:
+        script_path = path.dirname(path.abspath(__file__))
+        file_path = path.join(
+            script_path,
+            "lcmdiffusion",
+            "pipelines",
+            "latent_consistency_txt2img.py",
+        )
+        return file_path
+
+    def init(
+        self,
+        model_id: str,
+        use_openvino: bool = False,
+        device: str = "cpu",
+        use_local_model: bool = False,
+    ) -> None:
+        self.device = device
+        self.use_openvino = use_openvino
+        if self.pipeline is None or self.previous_model_id != model_id:
+            if self.use_openvino and DEVICE == "cpu":
+                if self.pipeline:
+                    del self.pipeline
+                scheduler = LCMScheduler.from_pretrained(
+                    model_id,
+                    subfolder="scheduler",
+                )
+                self.pipeline = OVLatentConsistencyModelPipeline.from_pretrained(
+                    model_id,
+                    scheduler=scheduler,
+                    compile=False,
+                    local_files_only=use_local_model,
+                )
+            else:
+                if self.pipeline:
+                    del self.pipeline
+
+                self.pipeline = DiffusionPipeline.from_pretrained(
+                    model_id,
+                    custom_pipeline=self._get_lcm_diffusion_pipeline_path(),
+                    custom_revision="main",
+                    local_files_only=use_local_model,
+                )
+                self.pipeline.to(
+                    torch_device=self.device,
+                    torch_dtype=torch.float32,
+                )
+            self.previous_model_id = model_id
+
+    def generate(
+        self,
+        lcm_diffusion_setting: LCMDiffusionSetting,
+        reshape: bool = False,
+    ) -> Any:
+        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 self.use_openvino and DEVICE == "cpu":
+            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
+
+        result_images = self.pipeline(
+            prompt=lcm_diffusion_setting.prompt,
+            num_inference_steps=lcm_diffusion_setting.inference_steps,
+            guidance_scale=lcm_diffusion_setting.guidance_scale,
+            width=lcm_diffusion_setting.image_width,
+            height=lcm_diffusion_setting.image_height,
+            output_type="pil",
+            num_images_per_prompt=lcm_diffusion_setting.number_of_images,
+        ).images
+
+        return result_images

backend/lcmdiffusion/pipelines/latent_consistency_txt2img.py

@@ -0,0 +1,730 @@
+# Copyright 2023 Stanford University Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
+# and https://github.com/hojonathanho/diffusion
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+
+from diffusers import AutoencoderKL, ConfigMixin, DiffusionPipeline, SchedulerMixin, UNet2DConditionModel, logging
+from diffusers.configuration_utils import register_to_config
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from diffusers.utils import BaseOutput
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class LatentConsistencyModelPipeline(DiffusionPipeline):
+    _optional_components = ["scheduler"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: "LCMScheduler",
+        safety_checker: StableDiffusionSafetyChecker,
+        feature_extractor: CLIPImageProcessor,
+        requires_safety_checker: bool = True,
+    ):
+        super().__init__()
+
+        scheduler = (
+            scheduler
+            if scheduler is not None
+            else LCMScheduler(
+                beta_start=0.00085, beta_end=0.0120, beta_schedule="scaled_linear", prediction_type="epsilon"
+            )
+        )
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+    def _encode_prompt(
+        self,
+        prompt,
+        device,
+        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
+            device: (`torch.device`):
+                torch device
+            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 is not None and isinstance(prompt, str):
+            pass
+        elif prompt is not None and isinstance(prompt, list):
+            len(prompt)
+        else:
+            prompt_embeds.shape[0]
+
+        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}"
+                )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = text_inputs.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            prompt_embeds = self.text_encoder(
+                text_input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            prompt_embeds = prompt_embeds[0]
+
+        if self.text_encoder is not None:
+            prompt_embeds_dtype = self.text_encoder.dtype
+        elif self.unet is not None:
+            prompt_embeds_dtype = self.unet.dtype
+        else:
+            prompt_embeds_dtype = prompt_embeds.dtype
+
+        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        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, device, 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").to(device)
+            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, device, 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).to(device)
+        else:
+            latents = latents.to(device)
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        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] = 768,
+        width: Optional[int] = 768,
+        guidance_scale: float = 7.5,
+        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,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        # 2. 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]
+
+        device = self._execution_device
+        # 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)
+
+        # 3. Encode input prompt
+        prompt_embeds = self._encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            prompt_embeds=prompt_embeds,
+        )
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, lcm_origin_steps)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variable
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            latents,
+        )
+        bs = batch_size * num_images_per_prompt
+
+        # 6. Get Guidance Scale Embedding
+        w = torch.tensor(guidance_scale).repeat(bs)
+        w_embedding = self.get_w_embedding(w, embedding_dim=256).to(device=device, dtype=latents.dtype)
+
+        # 7. LCM MultiStep Sampling Loop:
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                ts = torch.full((bs,), t, device=device, dtype=torch.long)
+                latents = latents.to(prompt_embeds.dtype)
+
+                # model prediction (v-prediction, eps, x)
+                model_pred = self.unet(
+                    latents,
+                    ts,
+                    timestep_cond=w_embedding,
+                    encoder_hidden_states=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents, denoised = self.scheduler.step(model_pred, i, t, latents, return_dict=False)
+
+                # # call the callback, if provided
+                # if i == len(timesteps) - 1:
+                progress_bar.update()
+
+        denoised = denoised.to(prompt_embeds.dtype)
+        if not output_type == "latent":
+            image = self.vae.decode(denoised / self.vae.config.scaling_factor, return_dict=False)[0]
+            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
+        else:
+            image = denoised
+            has_nsfw_concept = None
+
+        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]
+
+        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
+
+
+@dataclass
+# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM
+class LCMSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
+            `pred_original_sample` can be used to preview progress or for guidance.
+    """
+
+    prev_sample: torch.FloatTensor
+    denoised: Optional[torch.FloatTensor] = None
+
+
+# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
+def betas_for_alpha_bar(
+    num_diffusion_timesteps,
+    max_beta=0.999,
+    alpha_transform_type="cosine",
+):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
+    (1-beta) over time from t = [0,1].
+    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
+    to that part of the diffusion process.
+    Args:
+        num_diffusion_timesteps (`int`): the number of betas to produce.
+        max_beta (`float`): the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+        alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.
+                     Choose from `cosine` or `exp`
+    Returns:
+        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
+    """
+    if alpha_transform_type == "cosine":
+
+        def alpha_bar_fn(t):
+            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
+
+    elif alpha_transform_type == "exp":
+
+        def alpha_bar_fn(t):
+            return math.exp(t * -12.0)
+
+    else:
+        raise ValueError(f"Unsupported alpha_tranform_type: {alpha_transform_type}")
+
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
+    return torch.tensor(betas, dtype=torch.float32)
+
+
+def rescale_zero_terminal_snr(betas):
+    """
+    Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)
+    Args:
+        betas (`torch.FloatTensor`):
+            the betas that the scheduler is being initialized with.
+    Returns:
+        `torch.FloatTensor`: rescaled betas with zero terminal SNR
+    """
+    # Convert betas to alphas_bar_sqrt
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+    alphas_bar_sqrt = alphas_cumprod.sqrt()
+
+    # Store old values.
+    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
+    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
+
+    # Shift so the last timestep is zero.
+    alphas_bar_sqrt -= alphas_bar_sqrt_T
+
+    # Scale so the first timestep is back to the old value.
+    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+
+    # Convert alphas_bar_sqrt to betas
+    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
+    alphas = alphas_bar[1:] / alphas_bar[:-1]  # Revert cumprod
+    alphas = torch.cat([alphas_bar[0:1], alphas])
+    betas = 1 - alphas
+
+    return betas
+
+
+class LCMScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `LCMScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with
+    non-Markovian guidance.
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        beta_start (`float`, defaults to 0.0001):
+            The starting `beta` value of inference.
+        beta_end (`float`, defaults to 0.02):
+            The final `beta` value.
+        beta_schedule (`str`, defaults to `"linear"`):
+            The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
+            `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
+        trained_betas (`np.ndarray`, *optional*):
+            Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
+        clip_sample (`bool`, defaults to `True`):
+            Clip the predicted sample for numerical stability.
+        clip_sample_range (`float`, defaults to 1.0):
+            The maximum magnitude for sample clipping. Valid only when `clip_sample=True`.
+        set_alpha_to_one (`bool`, defaults to `True`):
+            Each diffusion step uses the alphas product value at that step and at the previous one. For the final step
+            there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`,
+            otherwise it uses the alpha value at step 0.
+        steps_offset (`int`, defaults to 0):
+            An offset added to the inference steps. You can use a combination of `offset=1` and
+            `set_alpha_to_one=False` to make the last step use step 0 for the previous alpha product like in Stable
+            Diffusion.
+        prediction_type (`str`, defaults to `epsilon`, *optional*):
+            Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
+            `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
+            Video](https://imagen.research.google/video/paper.pdf) paper).
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+            as Stable Diffusion.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True`.
+        timestep_spacing (`str`, defaults to `"leading"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        rescale_betas_zero_snr (`bool`, defaults to `False`):
+            Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and
+            dark samples instead of limiting it to samples with medium brightness. Loosely related to
+            [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).
+    """
+
+    # _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.0001,
+        beta_end: float = 0.02,
+        beta_schedule: str = "linear",
+        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        clip_sample: bool = True,
+        set_alpha_to_one: bool = True,
+        steps_offset: int = 0,
+        prediction_type: str = "epsilon",
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        clip_sample_range: float = 1.0,
+        sample_max_value: float = 1.0,
+        timestep_spacing: str = "leading",
+        rescale_betas_zero_snr: bool = False,
+    ):
+        if trained_betas is not None:
+            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
+        elif beta_schedule == "linear":
+            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+        elif beta_schedule == "scaled_linear":
+            # this schedule is very specific to the latent diffusion model.
+            self.betas = (
+                torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
+            )
+        elif beta_schedule == "squaredcos_cap_v2":
+            # Glide cosine schedule
+            self.betas = betas_for_alpha_bar(num_train_timesteps)
+        else:
+            raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")
+
+        # Rescale for zero SNR
+        if rescale_betas_zero_snr:
+            self.betas = rescale_zero_terminal_snr(self.betas)
+
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+
+        # At every step in ddim, we are looking into the previous alphas_cumprod
+        # For the final step, there is no previous alphas_cumprod because we are already at 0
+        # `set_alpha_to_one` decides whether we set this parameter simply to one or
+        # whether we use the final alpha of the "non-previous" one.
+        self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]
+
+        # standard deviation of the initial noise distribution
+        self.init_noise_sigma = 1.0
+
+        # setable values
+        self.num_inference_steps = None
+        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))
+
+    def scale_model_input(self, sample: torch.FloatTensor, timestep: Optional[int] = None) -> torch.FloatTensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+        Args:
+            sample (`torch.FloatTensor`):
+                The input sample.
+            timestep (`int`, *optional*):
+                The current timestep in the diffusion chain.
+        Returns:
+            `torch.FloatTensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    def _get_variance(self, timestep, prev_timestep):
+        alpha_prod_t = self.alphas_cumprod[timestep]
+        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
+        beta_prod_t = 1 - alpha_prod_t
+        beta_prod_t_prev = 1 - alpha_prod_t_prev
+
+        variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
+
+        return variance
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, height, width = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * height * width)
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+
+        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = torch.clamp(sample, -s, s) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, height, width)
+        sample = sample.to(dtype)
+
+        return sample
+
+    def set_timesteps(self, num_inference_steps: int, lcm_origin_steps: int, device: Union[str, torch.device] = None):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+        """
+
+        if num_inference_steps > self.config.num_train_timesteps:
+            raise ValueError(
+                f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"
+                f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
+                f" maximal {self.config.num_train_timesteps} timesteps."
+            )
+
+        self.num_inference_steps = num_inference_steps
+
+        # LCM Timesteps Setting:  # Linear Spacing
+        c = self.config.num_train_timesteps // lcm_origin_steps
+        lcm_origin_timesteps = np.asarray(list(range(1, lcm_origin_steps + 1))) * c - 1  # LCM Training  Steps Schedule
+        skipping_step = len(lcm_origin_timesteps) // num_inference_steps
+        timesteps = lcm_origin_timesteps[::-skipping_step][:num_inference_steps]  # LCM Inference Steps Schedule
+
+        self.timesteps = torch.from_numpy(timesteps.copy()).to(device)
+
+    def get_scalings_for_boundary_condition_discrete(self, t):
+        self.sigma_data = 0.5  # Default: 0.5
+
+        # By dividing 0.1: This is almost a delta function at t=0.
+        c_skip = self.sigma_data**2 / ((t / 0.1) ** 2 + self.sigma_data**2)
+        c_out = (t / 0.1) / ((t / 0.1) ** 2 + self.sigma_data**2) ** 0.5
+        return c_skip, c_out
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timeindex: int,
+        timestep: int,
+        sample: torch.FloatTensor,
+        eta: float = 0.0,
+        use_clipped_model_output: bool = False,
+        generator=None,
+        variance_noise: Optional[torch.FloatTensor] = None,
+        return_dict: bool = True,
+    ) -> Union[LCMSchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            eta (`float`):
+                The weight of noise for added noise in diffusion step.
+            use_clipped_model_output (`bool`, defaults to `False`):
+                If `True`, computes "corrected" `model_output` from the clipped predicted original sample. Necessary
+                because predicted original sample is clipped to [-1, 1] when `self.config.clip_sample` is `True`. If no
+                clipping has happened, "corrected" `model_output` would coincide with the one provided as input and
+                `use_clipped_model_output` has no effect.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            variance_noise (`torch.FloatTensor`):
+                Alternative to generating noise with `generator` by directly providing the noise for the variance
+                itself. Useful for methods such as [`CycleDiffusion`].
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`.
+        Returns:
+            [`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
+        """
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+            )
+
+        # 1. get previous step value
+        prev_timeindex = timeindex + 1
+        if prev_timeindex < len(self.timesteps):
+            prev_timestep = self.timesteps[prev_timeindex]
+        else:
+            prev_timestep = timestep
+
+        # 2. compute alphas, betas
+        alpha_prod_t = self.alphas_cumprod[timestep]
+        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
+
+        beta_prod_t = 1 - alpha_prod_t
+        beta_prod_t_prev = 1 - alpha_prod_t_prev
+
+        # 3. Get scalings for boundary conditions
+        c_skip, c_out = self.get_scalings_for_boundary_condition_discrete(timestep)
+
+        # 4. Different Parameterization:
+        parameterization = self.config.prediction_type
+
+        if parameterization == "epsilon":  # noise-prediction
+            pred_x0 = (sample - beta_prod_t.sqrt() * model_output) / alpha_prod_t.sqrt()
+
+        elif parameterization == "sample":  # x-prediction
+            pred_x0 = model_output
+
+        elif parameterization == "v_prediction":  # v-prediction
+            pred_x0 = alpha_prod_t.sqrt() * sample - beta_prod_t.sqrt() * model_output
+
+        # 4. Denoise model output using boundary conditions
+        denoised = c_out * pred_x0 + c_skip * sample
+
+        # 5. Sample z ~ N(0, I), For MultiStep Inference
+        # Noise is not used for one-step sampling.
+        if len(self.timesteps) > 1:
+            noise = torch.randn(model_output.shape).to(model_output.device)
+            prev_sample = alpha_prod_t_prev.sqrt() * denoised + beta_prod_t_prev.sqrt() * noise
+        else:
+            prev_sample = denoised
+
+        if not return_dict:
+            return (prev_sample, denoised)
+
+        return LCMSchedulerOutput(prev_sample=prev_sample, denoised=denoised)
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise
+    def add_noise(
+        self,
+        original_samples: torch.FloatTensor,
+        noise: torch.FloatTensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.FloatTensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
+        alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
+        timesteps = timesteps.to(original_samples.device)
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+        return noisy_samples
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity
+    def get_velocity(
+        self, sample: torch.FloatTensor, noise: torch.FloatTensor, timesteps: torch.IntTensor
+    ) -> torch.FloatTensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as sample
+        alphas_cumprod = self.alphas_cumprod.to(device=sample.device, dtype=sample.dtype)
+        timesteps = timesteps.to(sample.device)
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(sample.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
+        return velocity
+
+    def __len__(self):
+        return self.config.num_train_timesteps

backend/lcmdiffusion/pipelines/openvino/lcm_ov_pipeline.py

@@ -0,0 +1,390 @@
+# https://huggingface.co/deinferno/LCM_Dreamshaper_v7-openvino
+
+import inspect
+
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import List, Optional, Tuple, Union, Dict, Any, Callable, OrderedDict
+
+import numpy as np
+import openvino
+import torch
+
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from optimum.intel.openvino.modeling_diffusion import OVStableDiffusionPipeline, OVModelUnet, OVModelVaeDecoder, OVModelTextEncoder, OVModelVaeEncoder, VaeImageProcessor
+from optimum.utils import (
+    DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER,
+    DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER,
+    DIFFUSION_MODEL_UNET_SUBFOLDER,
+    DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER,
+    DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER,
+)
+
+
+from diffusers import logging
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+class LCMOVModelUnet(OVModelUnet):
+    def __call__(
+        self,
+        sample: np.ndarray,
+        timestep: np.ndarray,
+        encoder_hidden_states: np.ndarray,
+        timestep_cond: Optional[np.ndarray] = None,
+        text_embeds: Optional[np.ndarray] = None,
+        time_ids: Optional[np.ndarray] = None,
+    ):
+        self._compile()
+
+        inputs = {
+            "sample": sample,
+            "timestep": timestep,
+            "encoder_hidden_states": encoder_hidden_states,
+        }
+
+        if timestep_cond is not None:
+            inputs["timestep_cond"] = timestep_cond
+        if text_embeds is not None:
+            inputs["text_embeds"] = text_embeds
+        if time_ids is not None:
+            inputs["time_ids"] = time_ids
+
+        outputs = self.request(inputs, shared_memory=True)
+        return list(outputs.values())
+
+class OVLatentConsistencyModelPipeline(OVStableDiffusionPipeline):
+
+    def __init__(
+        self,
+        vae_decoder: openvino.runtime.Model,
+        text_encoder: openvino.runtime.Model,
+        unet: openvino.runtime.Model,
+        config: Dict[str, Any],
+        tokenizer: "CLIPTokenizer",
+        scheduler: Union["DDIMScheduler", "PNDMScheduler", "LMSDiscreteScheduler"],
+        feature_extractor: Optional["CLIPFeatureExtractor"] = None,
+        vae_encoder: Optional[openvino.runtime.Model] = None,
+        text_encoder_2: Optional[openvino.runtime.Model] = None,
+        tokenizer_2: Optional["CLIPTokenizer"] = None,
+        device: str = "CPU",
+        dynamic_shapes: bool = True,
+        compile: bool = True,
+        ov_config: Optional[Dict[str, str]] = None,
+        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
+        **kwargs,
+    ):
+        self._internal_dict = config
+        self._device = device.upper()
+        self.is_dynamic = dynamic_shapes
+        self.ov_config = ov_config if ov_config is not None else {}
+        self._model_save_dir = (
+            Path(model_save_dir.name) if isinstance(model_save_dir, TemporaryDirectory) else model_save_dir
+        )
+        self.vae_decoder = OVModelVaeDecoder(vae_decoder, self)
+        self.unet = LCMOVModelUnet(unet, self)
+        self.text_encoder = OVModelTextEncoder(text_encoder, self) if text_encoder is not None else None
+        self.text_encoder_2 = (
+            OVModelTextEncoder(text_encoder_2, self, model_name=DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER)
+            if text_encoder_2 is not None
+            else None
+        )
+        self.vae_encoder = OVModelVaeEncoder(vae_encoder, self) if vae_encoder is not None else None
+
+        if "block_out_channels" in self.vae_decoder.config:
+            self.vae_scale_factor = 2 ** (len(self.vae_decoder.config["block_out_channels"]) - 1)
+        else:
+            self.vae_scale_factor = 8
+
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+        self.tokenizer = tokenizer
+        self.tokenizer_2 = tokenizer_2
+        self.scheduler = scheduler
+        self.feature_extractor = feature_extractor
+        self.safety_checker = None
+        self.preprocessors = []
+
+        if self.is_dynamic:
+            self.reshape(batch_size=-1, height=-1, width=-1, num_images_per_prompt=-1)
+
+        if compile:
+            self.compile()
+
+        sub_models = {
+            DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER: self.text_encoder,
+            DIFFUSION_MODEL_UNET_SUBFOLDER: self.unet,
+            DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER: self.vae_decoder,
+            DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER: self.vae_encoder,
+            DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER: self.text_encoder_2,
+        }
+        for name in sub_models.keys():
+            self._internal_dict[name] = (
+                ("optimum", sub_models[name].__class__.__name__) if sub_models[name] is not None else (None, None)
+            )
+
+        self._internal_dict.pop("vae", None)
+
+    def _reshape_unet(
+        self,
+        model: openvino.runtime.Model,
+        batch_size: int = -1,
+        height: int = -1,
+        width: int = -1,
+        num_images_per_prompt: int = -1,
+        tokenizer_max_length: int = -1,
+    ):  
+        if batch_size == -1 or num_images_per_prompt == -1:
+            batch_size = -1
+        else:
+            batch_size = batch_size * num_images_per_prompt
+
+        height = height // self.vae_scale_factor if height > 0 else height
+        width = width // self.vae_scale_factor if width > 0 else width
+        shapes = {}
+        for inputs in model.inputs:
+            shapes[inputs] = inputs.get_partial_shape()
+            if inputs.get_any_name() == "timestep":
+                shapes[inputs][0] = 1
+            elif inputs.get_any_name() == "sample":
+                in_channels = self.unet.config.get("in_channels", None)
+                if in_channels is None:
+                    in_channels = shapes[inputs][1]
+                    if in_channels.is_dynamic:
+                        logger.warning(
+                            "Could not identify `in_channels` from the unet configuration, to statically reshape the unet please provide a configuration."
+                        ) 
+                        self.is_dynamic = True
+    
+                shapes[inputs] = [batch_size, in_channels, height, width]
+            elif inputs.get_any_name() == "timestep_cond":
+                shapes[inputs] = [batch_size, inputs.get_partial_shape()[1]]
+            elif inputs.get_any_name() == "text_embeds":
+                shapes[inputs] = [batch_size, self.text_encoder_2.config["projection_dim"]]
+            elif inputs.get_any_name() == "time_ids":
+                shapes[inputs] = [batch_size, inputs.get_partial_shape()[1]]
+            else:
+                shapes[inputs][0] = batch_size
+                shapes[inputs][1] = tokenizer_max_length
+        model.reshape(shapes)
+        return model
+
+    def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=np.float32):
+        """
+        see https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
+        Args:
+        timesteps: np.array: 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.
+
+        half_dim = embedding_dim // 2
+        emb = np.log(np.array(10000.)) / (half_dim - 1)
+        emb = np.exp(np.arange(half_dim, dtype=dtype) * -emb)
+        emb = w.astype(dtype)[:, None] * emb[None, :]
+        emb = np.concatenate([np.sin(emb), np.cos(emb)], axis=1)
+        if embedding_dim % 2 == 1:  # zero pad
+            emb = np.pad(emb, (0, 1))
+        assert emb.shape == (w.shape[0], embedding_dim)
+        return emb
+
+    # Adapted from https://github.com/huggingface/optimum/blob/15b8d1eed4d83c5004d3b60f6b6f13744b358f01/optimum/pipelines/diffusers/pipeline_stable_diffusion.py#L201
+    def __call__(
+        self,
+        prompt: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 4,
+        original_inference_steps: int = None,
+        guidance_scale: float = 7.5,
+        num_images_per_prompt: int = 1,
+        eta: float = 0.0,
+        generator: Optional[np.random.RandomState] = None,
+        latents: Optional[np.ndarray] = None,
+        prompt_embeds: Optional[np.ndarray] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
+        callback_steps: int = 1,
+        guidance_rescale: float = 0.0,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`Optional[Union[str, List[str]]]`, defaults to None):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            height (`Optional[int]`, defaults to None):
+                The height in pixels of the generated image.
+            width (`Optional[int]`, defaults to None):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, defaults to 4):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            original_inference_steps (`int`, *optional*):
+                The original number of inference steps use to generate a linearly-spaced timestep schedule, from which
+                we will draw `num_inference_steps` evenly spaced timesteps from as our final timestep schedule,
+                following the Skipping-Step method in the paper (see Section 4.3). If not set this will default to the
+                scheduler's `original_inference_steps` attribute.
+            guidance_scale (`float`, defaults to 7.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            num_images_per_prompt (`int`, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`Optional[np.random.RandomState]`, defaults to `None`)::
+                A np.random.RandomState to make generation deterministic.
+            latents (`Optional[np.ndarray]`, defaults to `None`):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`Optional[np.ndarray]`, defaults to `None`):
+                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.
+            output_type (`str`, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
+                plain tuple.
+            callback (Optional[Callable], defaults to `None`):
+                A function that will be called every `callback_steps` steps during inference. The function will be
+                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, defaults to 1):
+                The frequency at which the `callback` function will be called. If not specified, the callback will be
+                called at every step.
+            guidance_rescale (`float`, defaults to 0.0):
+                Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
+                Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
+                [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
+                Guidance rescale factor should fix overexposure when using zero terminal SNR.
+
+        Returns:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
+            When returning a tuple, the first element is a list with the generated images, and the second element is a
+            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
+            (nsfw) content, according to the `safety_checker`.
+        """
+        height = height or self.unet.config.get("sample_size", 64) * self.vae_scale_factor
+        width = width or self.unet.config.get("sample_size", 64) * self.vae_scale_factor
+
+        # check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt, height, width, callback_steps, None, prompt_embeds, None
+        )
+
+        # define call parameters
+        if isinstance(prompt, str):
+            batch_size = 1
+        elif isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if generator is None:
+            generator = np.random
+
+        # Create torch.Generator instance with same state as np.random.RandomState
+        torch_generator = torch.Generator().manual_seed(int(generator.get_state()[1][0]))
+
+        #do_classifier_free_guidance = guidance_scale > 1.0
+
+        # NOTE: when a LCM is distilled from an LDM via latent consistency distillation (Algorithm 1) with guided
+        # distillation, the forward pass of the LCM learns to approximate sampling from the LDM using CFG with the
+        # unconditional prompt "" (the empty string). Due to this, LCMs currently do not support negative prompts.
+        prompt_embeds = self._encode_prompt(
+            prompt,
+            num_images_per_prompt,
+            False,
+            negative_prompt=None,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=None,
+        )
+
+        # set timesteps
+        self.scheduler.set_timesteps(num_inference_steps, "cpu", original_inference_steps=original_inference_steps)
+        timesteps = self.scheduler.timesteps
+
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            self.unet.config.get("in_channels", 4),
+            height,
+            width,
+            prompt_embeds.dtype,
+            generator,
+            latents,
+        )
+
+        # Get Guidance Scale Embedding
+        w = np.tile(guidance_scale - 1, batch_size * num_images_per_prompt)
+        w_embedding = self.get_guidance_scale_embedding(w, embedding_dim=self.unet.config.get("time_cond_proj_dim", 256))
+
+        # Adapted from diffusers to extend it for other runtimes than ORT
+        timestep_dtype = self.unet.input_dtype.get("timestep", np.float32)
+
+        # 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(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = torch_generator
+
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        for i, t in enumerate(self.progress_bar(timesteps)):
+
+            # predict the noise residual
+            timestep = np.array([t], dtype=timestep_dtype)
+            
+            noise_pred = self.unet(sample=latents, timestep=timestep, timestep_cond = w_embedding, encoder_hidden_states=prompt_embeds)[0]
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents, denoised = self.scheduler.step(
+                torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs, return_dict = False
+            )
+
+            latents, denoised = latents.numpy(), denoised.numpy()
+
+            # call the callback, if provided
+            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                if callback is not None and i % callback_steps == 0:
+                    callback(i, t, latents)
+
+        if output_type == "latent":
+            image = latents
+            has_nsfw_concept = None
+        else:
+            denoised /= self.vae_decoder.config.get("scaling_factor", 0.18215)
+            # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1
+            image = np.concatenate(
+                [self.vae_decoder(latent_sample=denoised[i : i + 1])[0] for i in range(latents.shape[0])]
+            )
+            image, has_nsfw_concept = self.run_safety_checker(image)
+
+        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]
+
+        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)

backend/lcmdiffusion/pipelines/openvino/lcm_scheduler.py

@@ -0,0 +1,529 @@
+# Copyright 2023 Stanford University Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
+# and https://github.com/hojonathanho/diffusion
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput, logging
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class LCMSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
+            `pred_original_sample` can be used to preview progress or for guidance.
+    """
+
+    prev_sample: torch.FloatTensor
+    denoised: Optional[torch.FloatTensor] = None
+
+
+# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
+def betas_for_alpha_bar(
+    num_diffusion_timesteps,
+    max_beta=0.999,
+    alpha_transform_type="cosine",
+):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
+    (1-beta) over time from t = [0,1].
+
+    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
+    to that part of the diffusion process.
+
+
+    Args:
+        num_diffusion_timesteps (`int`): the number of betas to produce.
+        max_beta (`float`): the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+        alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.
+                     Choose from `cosine` or `exp`
+
+    Returns:
+        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
+    """
+    if alpha_transform_type == "cosine":
+
+        def alpha_bar_fn(t):
+            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
+
+    elif alpha_transform_type == "exp":
+
+        def alpha_bar_fn(t):
+            return math.exp(t * -12.0)
+
+    else:
+        raise ValueError(f"Unsupported alpha_tranform_type: {alpha_transform_type}")
+
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
+    return torch.tensor(betas, dtype=torch.float32)
+
+
+# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr
+def rescale_zero_terminal_snr(betas: torch.FloatTensor) -> torch.FloatTensor:
+    """
+    Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)
+
+
+    Args:
+        betas (`torch.FloatTensor`):
+            the betas that the scheduler is being initialized with.
+
+    Returns:
+        `torch.FloatTensor`: rescaled betas with zero terminal SNR
+    """
+    # Convert betas to alphas_bar_sqrt
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+    alphas_bar_sqrt = alphas_cumprod.sqrt()
+
+    # Store old values.
+    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
+    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
+
+    # Shift so the last timestep is zero.
+    alphas_bar_sqrt -= alphas_bar_sqrt_T
+
+    # Scale so the first timestep is back to the old value.
+    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+
+    # Convert alphas_bar_sqrt to betas
+    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
+    alphas = alphas_bar[1:] / alphas_bar[:-1]  # Revert cumprod
+    alphas = torch.cat([alphas_bar[0:1], alphas])
+    betas = 1 - alphas
+
+    return betas
+
+
+class LCMScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `LCMScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with
+    non-Markovian guidance.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. [`~ConfigMixin`] takes care of storing all config
+    attributes that are passed in the scheduler's `__init__` function, such as `num_train_timesteps`. They can be
+    accessed via `scheduler.config.num_train_timesteps`. [`SchedulerMixin`] provides general loading and saving
+    functionality via the [`SchedulerMixin.save_pretrained`] and [`~SchedulerMixin.from_pretrained`] functions.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        beta_start (`float`, defaults to 0.0001):
+            The starting `beta` value of inference.
+        beta_end (`float`, defaults to 0.02):
+            The final `beta` value.
+        beta_schedule (`str`, defaults to `"linear"`):
+            The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
+            `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
+        trained_betas (`np.ndarray`, *optional*):
+            Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
+        original_inference_steps (`int`, *optional*, defaults to 50):
+            The default number of inference steps used to generate a linearly-spaced timestep schedule, from which we
+            will ultimately take `num_inference_steps` evenly spaced timesteps to form the final timestep schedule.
+        clip_sample (`bool`, defaults to `True`):
+            Clip the predicted sample for numerical stability.
+        clip_sample_range (`float`, defaults to 1.0):
+            The maximum magnitude for sample clipping. Valid only when `clip_sample=True`.
+        set_alpha_to_one (`bool`, defaults to `True`):
+            Each diffusion step uses the alphas product value at that step and at the previous one. For the final step
+            there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`,
+            otherwise it uses the alpha value at step 0.
+        steps_offset (`int`, defaults to 0):
+            An offset added to the inference steps. You can use a combination of `offset=1` and
+            `set_alpha_to_one=False` to make the last step use step 0 for the previous alpha product like in Stable
+            Diffusion.
+        prediction_type (`str`, defaults to `epsilon`, *optional*):
+            Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
+            `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
+            Video](https://imagen.research.google/video/paper.pdf) paper).
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+            as Stable Diffusion.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True`.
+        timestep_spacing (`str`, defaults to `"leading"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        rescale_betas_zero_snr (`bool`, defaults to `False`):
+            Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and
+            dark samples instead of limiting it to samples with medium brightness. Loosely related to
+            [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).
+    """
+
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.00085,
+        beta_end: float = 0.012,
+        beta_schedule: str = "scaled_linear",
+        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        original_inference_steps: int = 50,
+        clip_sample: bool = False,
+        clip_sample_range: float = 1.0,
+        set_alpha_to_one: bool = True,
+        steps_offset: int = 0,
+        prediction_type: str = "epsilon",
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        sample_max_value: float = 1.0,
+        timestep_spacing: str = "leading",
+        rescale_betas_zero_snr: bool = False,
+    ):
+        if trained_betas is not None:
+            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
+        elif beta_schedule == "linear":
+            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+        elif beta_schedule == "scaled_linear":
+            # this schedule is very specific to the latent diffusion model.
+            self.betas = (
+                torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
+            )
+        elif beta_schedule == "squaredcos_cap_v2":
+            # Glide cosine schedule
+            self.betas = betas_for_alpha_bar(num_train_timesteps)
+        else:
+            raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")
+
+        # Rescale for zero SNR
+        if rescale_betas_zero_snr:
+            self.betas = rescale_zero_terminal_snr(self.betas)
+
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+
+        # At every step in ddim, we are looking into the previous alphas_cumprod
+        # For the final step, there is no previous alphas_cumprod because we are already at 0
+        # `set_alpha_to_one` decides whether we set this parameter simply to one or
+        # whether we use the final alpha of the "non-previous" one.
+        self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]
+
+        # standard deviation of the initial noise distribution
+        self.init_noise_sigma = 1.0
+
+        # setable values
+        self.num_inference_steps = None
+        self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))
+
+        self._step_index = None
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
+    def _init_step_index(self, timestep):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.to(self.timesteps.device)
+
+        index_candidates = (self.timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        if len(index_candidates) > 1:
+            step_index = index_candidates[1]
+        else:
+            step_index = index_candidates[0]
+
+        self._step_index = step_index.item()
+
+    @property
+    def step_index(self):
+        return self._step_index
+
+    def scale_model_input(self, sample: torch.FloatTensor, timestep: Optional[int] = None) -> torch.FloatTensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The input sample.
+            timestep (`int`, *optional*):
+                The current timestep in the diffusion chain.
+        Returns:
+            `torch.FloatTensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, *remaining_dims = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = torch.clamp(sample, -s, s) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
+        sample = sample.to(dtype)
+
+        return sample
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        device: Union[str, torch.device] = None,
+        original_inference_steps: Optional[int] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+            original_inference_steps (`int`, *optional*):
+                The original number of inference steps, which will be used to generate a linearly-spaced timestep
+                schedule (which is different from the standard `diffusers` implementation). We will then take
+                `num_inference_steps` timesteps from this schedule, evenly spaced in terms of indices, and use that as
+                our final timestep schedule. If not set, this will default to the `original_inference_steps` attribute.
+        """
+
+        if num_inference_steps > self.config.num_train_timesteps:
+            raise ValueError(
+                f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"
+                f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
+                f" maximal {self.config.num_train_timesteps} timesteps."
+            )
+
+        self.num_inference_steps = num_inference_steps
+        original_steps = (
+            original_inference_steps if original_inference_steps is not None else self.original_inference_steps
+        )
+
+        if original_steps > self.config.num_train_timesteps:
+            raise ValueError(
+                f"`original_steps`: {original_steps} cannot be larger than `self.config.train_timesteps`:"
+                f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
+                f" maximal {self.config.num_train_timesteps} timesteps."
+            )
+
+        if num_inference_steps > original_steps:
+            raise ValueError(
+                f"`num_inference_steps`: {num_inference_steps} cannot be larger than `original_inference_steps`:"
+                f" {original_steps} because the final timestep schedule will be a subset of the"
+                f" `original_inference_steps`-sized initial timestep schedule."
+            )
+
+        # LCM Timesteps Setting
+        # Currently, only linear spacing is supported.
+        c = self.config.num_train_timesteps // original_steps
+        # LCM Training Steps Schedule
+        lcm_origin_timesteps = np.asarray(list(range(1, original_steps + 1))) * c - 1
+        skipping_step = len(lcm_origin_timesteps) // num_inference_steps
+        # LCM Inference Steps Schedule
+        timesteps = lcm_origin_timesteps[::-skipping_step][:num_inference_steps]
+
+        self.timesteps = torch.from_numpy(timesteps.copy()).to(device=device, dtype=torch.long)
+
+        self._step_index = None
+
+    def get_scalings_for_boundary_condition_discrete(self, t):
+        self.sigma_data = 0.5  # Default: 0.5
+
+        # By dividing 0.1: This is almost a delta function at t=0.
+        c_skip = self.sigma_data**2 / ((t / 0.1) ** 2 + self.sigma_data**2)
+        c_out = (t / 0.1) / ((t / 0.1) ** 2 + self.sigma_data**2) ** 0.5
+        return c_skip, c_out
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        sample: torch.FloatTensor,
+        generator: Optional[torch.Generator] = None,
+        return_dict: bool = True,
+    ) -> Union[LCMSchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`.
+        Returns:
+            [`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
+        """
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # 1. get previous step value
+        prev_step_index = self.step_index + 1
+        if prev_step_index < len(self.timesteps):
+            prev_timestep = self.timesteps[prev_step_index]
+        else:
+            prev_timestep = timestep
+
+        # 2. compute alphas, betas
+        alpha_prod_t = self.alphas_cumprod[timestep]
+        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
+
+        beta_prod_t = 1 - alpha_prod_t
+        beta_prod_t_prev = 1 - alpha_prod_t_prev
+
+        # 3. Get scalings for boundary conditions
+        c_skip, c_out = self.get_scalings_for_boundary_condition_discrete(timestep)
+
+        # 4. Compute the predicted original sample x_0 based on the model parameterization
+        if self.config.prediction_type == "epsilon":  # noise-prediction
+            predicted_original_sample = (sample - beta_prod_t.sqrt() * model_output) / alpha_prod_t.sqrt()
+        elif self.config.prediction_type == "sample":  # x-prediction
+            predicted_original_sample = model_output
+        elif self.config.prediction_type == "v_prediction":  # v-prediction
+            predicted_original_sample = alpha_prod_t.sqrt() * sample - beta_prod_t.sqrt() * model_output
+        else:
+            raise ValueError(
+                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or"
+                " `v_prediction` for `LCMScheduler`."
+            )
+
+        # 5. Clip or threshold "predicted x_0"
+        if self.config.thresholding:
+            predicted_original_sample = self._threshold_sample(predicted_original_sample)
+        elif self.config.clip_sample:
+            predicted_original_sample = predicted_original_sample.clamp(
+                -self.config.clip_sample_range, self.config.clip_sample_range
+            )
+
+        # 6. Denoise model output using boundary conditions
+        denoised = c_out * predicted_original_sample + c_skip * sample
+
+        # 7. Sample and inject noise z ~ N(0, I) for MultiStep Inference
+        # Noise is not used for one-step sampling.
+        if len(self.timesteps) > 1:
+            noise = randn_tensor(model_output.shape, generator=generator, device=model_output.device)
+            prev_sample = alpha_prod_t_prev.sqrt() * denoised + beta_prod_t_prev.sqrt() * noise
+        else:
+            prev_sample = denoised
+
+        # upon completion increase step index by one
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample, denoised)
+
+        return LCMSchedulerOutput(prev_sample=prev_sample, denoised=denoised)
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise
+    def add_noise(
+        self,
+        original_samples: torch.FloatTensor,
+        noise: torch.FloatTensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.FloatTensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
+        alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
+        timesteps = timesteps.to(original_samples.device)
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+        return noisy_samples
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity
+    def get_velocity(
+        self, sample: torch.FloatTensor, noise: torch.FloatTensor, timesteps: torch.IntTensor
+    ) -> torch.FloatTensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as sample
+        alphas_cumprod = self.alphas_cumprod.to(device=sample.device, dtype=sample.dtype)
+        timesteps = timesteps.to(sample.device)
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(sample.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
+        return velocity
+
+    def __len__(self):
+        return self.config.num_train_timesteps

backend/models/lcmdiffusion_setting.py

@@ -0,0 +1,19 @@
+from typing import Optional
+
+from pydantic import BaseModel
+from constants import LCM_DEFAULT_MODEL
+
+
+class LCMDiffusionSetting(BaseModel):
+    lcm_model_id: str = LCM_DEFAULT_MODEL
+    prompt: str = ""
+    image_height: Optional[int] = 512
+    image_width: Optional[int] = 512
+    inference_steps: Optional[int] = 4
+    guidance_scale: Optional[float] = 8
+    number_of_images: Optional[int] = 1
+    seed: Optional[int] = -1
+    use_openvino: bool = False
+    use_seed: bool = False
+    use_offline_model: bool = False
+    use_safety_checker: bool = True

constants.py

@@ -0,0 +1,10 @@
+from os import environ
+
+APP_VERSION = "v1.0.0 beta 7"
+LCM_DEFAULT_MODEL = "SimianLuo/LCM_Dreamshaper_v7"
+LCM_DEFAULT_MODEL_OPENVINO = "deinferno/LCM_Dreamshaper_v7-openvino"
+APP_NAME = "FastSD CPU"
+APP_SETTINGS_FILE = "settings.yaml"
+RESULTS_DIRECTORY = "results"
+CONFIG_DIRECTORY = "configs"
+DEVICE = environ.get("DEVICE", "cpu")

context.py

@@ -0,0 +1,44 @@
+from typing import Any
+from app_settings import Settings
+from models.interface_types import InterfaceType
+from backend.lcm_text_to_image import LCMTextToImage
+from time import time
+from backend.image_saver import ImageSaver
+from pprint import pprint
+
+
+class Context:
+    def __init__(
+        self,
+        interface_type: InterfaceType,
+        device="cpu",
+    ):
+        self.interface_type = interface_type
+        self.lcm_text_to_image = LCMTextToImage(device)
+
+    def generate_text_to_image(
+        self,
+        settings: Settings,
+        reshape: bool = False,
+        device: str = "cpu",
+    ) -> Any:
+        tick = time()
+        pprint(settings.lcm_diffusion_setting.model_dump())
+        self.lcm_text_to_image.init(
+            settings.lcm_diffusion_setting.lcm_model_id,
+            settings.lcm_diffusion_setting.use_openvino,
+            device,
+            settings.lcm_diffusion_setting.use_offline_model,
+        )
+        images = self.lcm_text_to_image.generate(
+            settings.lcm_diffusion_setting,
+            reshape,
+        )
+        elapsed = time() - tick
+        ImageSaver.save_images(
+            settings.results_path,
+            images=images,
+            lcm_diffusion_setting=settings.lcm_diffusion_setting,
+        )
+        print(f"Elapsed time : {elapsed:.2f} seconds")
+        return images

frontend/gui/app_window.py

@@ -0,0 +1,435 @@
+from PyQt5.QtWidgets import (
+    QWidget,
+    QPushButton,
+    QHBoxLayout,
+    QVBoxLayout,
+    QLabel,
+    QLineEdit,
+    QMainWindow,
+    QSlider,
+    QTabWidget,
+    QSpacerItem,
+    QSizePolicy,
+    QComboBox,
+    QCheckBox,
+    QTextEdit,
+    QToolButton,
+    QFileDialog,
+)
+
+from PyQt5.QtGui import QPixmap, QDesktopServices
+from PyQt5.QtCore import QSize, QThreadPool, Qt, QUrl
+
+from PIL.ImageQt import ImageQt
+from constants import (
+    LCM_DEFAULT_MODEL,
+    LCM_DEFAULT_MODEL_OPENVINO,
+    APP_NAME,
+    APP_VERSION,
+)
+from frontend.gui.image_generator_worker import ImageGeneratorWorker
+from app_settings import AppSettings
+from paths import FastStableDiffusionPaths
+from frontend.utils import is_reshape_required
+from context import Context
+from models.interface_types import InterfaceType
+from constants import DEVICE
+from frontend.utils import enable_openvino_controls
+
+
+class MainWindow(QMainWindow):
+    def __init__(self, config: AppSettings):
+        super().__init__()
+        self.setWindowTitle(APP_NAME)
+        self.setFixedSize(QSize(600, 620))
+        self.init_ui()
+        self.pipeline = None
+        self.threadpool = QThreadPool()
+        self.config = config
+        self.device = "cpu"
+        self.previous_width = 0
+        self.previous_height = 0
+        self.previous_model = ""
+        self.previous_num_of_images = 0
+        self.context = Context(InterfaceType.GUI)
+        self.init_ui_values()
+        self.gen_images = []
+        self.image_index = 0
+        print(f"Output path : {  self.config.settings.results_path}")
+
+    def init_ui_values(self):
+        self.lcm_model.setEnabled(
+            not self.config.settings.lcm_diffusion_setting.use_openvino
+        )
+        self.guidance.setValue(
+            int(self.config.settings.lcm_diffusion_setting.guidance_scale * 10)
+        )
+        self.seed_value.setEnabled(self.config.settings.lcm_diffusion_setting.use_seed)
+        self.safety_checker.setChecked(
+            self.config.settings.lcm_diffusion_setting.use_safety_checker
+        )
+        self.use_openvino_check.setChecked(
+            self.config.settings.lcm_diffusion_setting.use_openvino
+        )
+        self.width.setCurrentText(
+            str(self.config.settings.lcm_diffusion_setting.image_width)
+        )
+        self.height.setCurrentText(
+            str(self.config.settings.lcm_diffusion_setting.image_height)
+        )
+        self.inference_steps.setValue(
+            int(self.config.settings.lcm_diffusion_setting.inference_steps)
+        )
+        self.seed_check.setChecked(self.config.settings.lcm_diffusion_setting.use_seed)
+        self.seed_value.setText(str(self.config.settings.lcm_diffusion_setting.seed))
+        self.use_local_model_folder.setChecked(
+            self.config.settings.lcm_diffusion_setting.use_offline_model
+        )
+        self.results_path.setText(self.config.settings.results_path)
+        self.num_images.setValue(
+            self.config.settings.lcm_diffusion_setting.number_of_images
+        )
+
+    def init_ui(self):
+        self.create_main_tab()
+        self.create_settings_tab()
+        self.create_about_tab()
+        self.show()
+
+    def create_main_tab(self):
+        self.img = QLabel("<<Image>>")
+        self.img.setAlignment(Qt.AlignCenter)
+        self.img.setFixedSize(QSize(512, 512))
+
+        self.prompt = QTextEdit()
+        self.prompt.setPlaceholderText("A fantasy landscape")
+        self.prompt.setAcceptRichText(False)
+        self.generate = QPushButton("Generate")
+        self.generate.clicked.connect(self.text_to_image)
+        self.prompt.setFixedHeight(35)
+        self.browse_results = QPushButton("...")
+        self.browse_results.setFixedWidth(30)
+        self.browse_results.clicked.connect(self.on_open_results_folder)
+        self.browse_results.setToolTip("Open output folder")
+
+        hlayout = QHBoxLayout()
+        hlayout.addWidget(self.prompt)
+        hlayout.addWidget(self.generate)
+        hlayout.addWidget(self.browse_results)
+
+        self.previous_img_btn = QToolButton()
+        self.previous_img_btn.setText("<")
+        self.previous_img_btn.clicked.connect(self.on_show_previous_image)
+        self.next_img_btn = QToolButton()
+        self.next_img_btn.setText(">")
+        self.next_img_btn.clicked.connect(self.on_show_next_image)
+        hlayout_nav = QHBoxLayout()
+        hlayout_nav.addWidget(self.previous_img_btn)
+        hlayout_nav.addWidget(self.img)
+        hlayout_nav.addWidget(self.next_img_btn)
+
+        vlayout = QVBoxLayout()
+        vlayout.addLayout(hlayout_nav)
+        vlayout.addLayout(hlayout)
+
+        self.tab_widget = QTabWidget(self)
+        self.tab_main = QWidget()
+        self.tab_settings = QWidget()
+        self.tab_about = QWidget()
+        self.tab_main.setLayout(vlayout)
+
+        self.tab_widget.addTab(self.tab_main, "Text to Image")
+        self.tab_widget.addTab(self.tab_settings, "Settings")
+        self.tab_widget.addTab(self.tab_about, "About")
+
+        self.setCentralWidget(self.tab_widget)
+        self.use_seed = False
+
+    def create_settings_tab(self):
+        model_hlayout = QHBoxLayout()
+        self.lcm_model_label = QLabel("Latent Consistency Model:")
+        self.lcm_model = QLineEdit(LCM_DEFAULT_MODEL)
+        model_hlayout.addWidget(self.lcm_model_label)
+        model_hlayout.addWidget(self.lcm_model)
+
+        self.inference_steps_value = QLabel("Number of inference steps: 4")
+        self.inference_steps = QSlider(orientation=Qt.Orientation.Horizontal)
+        self.inference_steps.setMaximum(25)
+        self.inference_steps.setMinimum(1)
+        self.inference_steps.setValue(4)
+        self.inference_steps.valueChanged.connect(self.update_steps_label)
+
+        self.num_images_value = QLabel("Number of images: 1")
+        self.num_images = QSlider(orientation=Qt.Orientation.Horizontal)
+        self.num_images.setMaximum(100)
+        self.num_images.setMinimum(1)
+        self.num_images.setValue(1)
+        self.num_images.valueChanged.connect(self.update_num_images_label)
+
+        self.guidance_value = QLabel("Guidance scale: 8")
+        self.guidance = QSlider(orientation=Qt.Orientation.Horizontal)
+        self.guidance.setMaximum(200)
+        self.guidance.setMinimum(10)
+        self.guidance.setValue(80)
+        self.guidance.valueChanged.connect(self.update_guidance_label)
+
+        self.width_value = QLabel("Width :")
+        self.width = QComboBox(self)
+        self.width.addItem("256")
+        self.width.addItem("512")
+        self.width.addItem("768")
+        self.width.setCurrentText("512")
+        self.width.currentIndexChanged.connect(self.on_width_changed)
+
+        self.height_value = QLabel("Height :")
+        self.height = QComboBox(self)
+        self.height.addItem("256")
+        self.height.addItem("512")
+        self.height.addItem("768")
+        self.height.setCurrentText("512")
+        self.height.currentIndexChanged.connect(self.on_height_changed)
+
+        self.seed_check = QCheckBox("Use seed")
+        self.seed_value = QLineEdit()
+        self.seed_value.setInputMask("9999999999")
+        self.seed_value.setText("123123")
+        self.seed_check.stateChanged.connect(self.seed_changed)
+
+        self.safety_checker = QCheckBox("Use safety checker")
+        self.safety_checker.setChecked(True)
+        self.safety_checker.stateChanged.connect(self.use_safety_checker_changed)
+
+        self.use_openvino_check = QCheckBox("Use OpenVINO")
+        self.use_openvino_check.setChecked(False)
+        self.use_local_model_folder = QCheckBox(
+            "Use locally cached model or downloaded model folder(offline)"
+        )
+        self.use_openvino_check.setEnabled(enable_openvino_controls())
+        self.use_local_model_folder.setChecked(False)
+        self.use_local_model_folder.stateChanged.connect(self.use_offline_model_changed)
+        self.use_openvino_check.stateChanged.connect(self.use_openvino_changed)
+
+        hlayout = QHBoxLayout()
+        hlayout.addWidget(self.seed_check)
+        hlayout.addWidget(self.seed_value)
+        hspacer = QSpacerItem(20, 10, QSizePolicy.Expanding, QSizePolicy.Minimum)
+        slider_hspacer = QSpacerItem(20, 10, QSizePolicy.Expanding, QSizePolicy.Minimum)
+
+        self.results_path_label = QLabel("Output path:")
+        self.results_path = QLineEdit()
+        self.results_path.textChanged.connect(self.on_path_changed)
+        self.browse_folder_btn = QToolButton()
+        self.browse_folder_btn.setText("...")
+        self.browse_folder_btn.clicked.connect(self.on_browse_folder)
+
+        self.reset = QPushButton("Reset All")
+        self.reset.clicked.connect(self.reset_all_settings)
+
+        vlayout = QVBoxLayout()
+        vspacer = QSpacerItem(20, 20, QSizePolicy.Minimum, QSizePolicy.Expanding)
+        vlayout.addItem(hspacer)
+        vlayout.addLayout(model_hlayout)
+        vlayout.addWidget(self.use_local_model_folder)
+        vlayout.addItem(slider_hspacer)
+        vlayout.addWidget(self.inference_steps_value)
+        vlayout.addWidget(self.inference_steps)
+        vlayout.addWidget(self.num_images_value)
+        vlayout.addWidget(self.num_images)
+        vlayout.addWidget(self.width_value)
+        vlayout.addWidget(self.width)
+        vlayout.addWidget(self.height_value)
+        vlayout.addWidget(self.height)
+        vlayout.addWidget(self.guidance_value)
+        vlayout.addWidget(self.guidance)
+        vlayout.addLayout(hlayout)
+        vlayout.addWidget(self.safety_checker)
+        vlayout.addWidget(self.use_openvino_check)
+        vlayout.addWidget(self.results_path_label)
+        hlayout_path = QHBoxLayout()
+        hlayout_path.addWidget(self.results_path)
+        hlayout_path.addWidget(self.browse_folder_btn)
+        vlayout.addLayout(hlayout_path)
+        self.tab_settings.setLayout(vlayout)
+        hlayout_reset = QHBoxLayout()
+        hspacer = QSpacerItem(20, 20, QSizePolicy.Expanding, QSizePolicy.Minimum)
+        hlayout_reset.addItem(hspacer)
+        hlayout_reset.addWidget(self.reset)
+        vlayout.addLayout(hlayout_reset)
+        vlayout.addItem(vspacer)
+
+    def create_about_tab(self):
+        self.label = QLabel()
+        self.label.setAlignment(Qt.AlignCenter)
+        self.label.setText(
+            f"""<h1>FastSD CPU {APP_VERSION}</h1> 
+               <h3>(c)2023 - Rupesh Sreeraman</h3>
+                <h3>Faster stable diffusion on CPU</h3>
+                 <h3>Based on Latent Consistency Models</h3>
+                <h3>GitHub : https://github.com/rupeshs/fastsdcpu/</h3>"""
+        )
+
+        vlayout = QVBoxLayout()
+        vlayout.addWidget(self.label)
+        self.tab_about.setLayout(vlayout)
+
+    def on_show_next_image(self):
+        if self.image_index != len(self.gen_images) - 1 and len(self.gen_images) > 0:
+            self.previous_img_btn.setEnabled(True)
+            self.image_index += 1
+            self.img.setPixmap(self.gen_images[self.image_index])
+            if self.image_index == len(self.gen_images) - 1:
+                self.next_img_btn.setEnabled(False)
+
+    def on_open_results_folder(self):
+        QDesktopServices.openUrl(QUrl.fromLocalFile(self.config.settings.results_path))
+
+    def on_show_previous_image(self):
+        if self.image_index != 0:
+            self.next_img_btn.setEnabled(True)
+            self.image_index -= 1
+            self.img.setPixmap(self.gen_images[self.image_index])
+            if self.image_index == 0:
+                self.previous_img_btn.setEnabled(False)
+
+    def on_path_changed(self, text):
+        self.config.settings.results_path = text
+
+    def on_browse_folder(self):
+        options = QFileDialog.Options()
+        options |= QFileDialog.ShowDirsOnly
+
+        folder_path = QFileDialog.getExistingDirectory(
+            self, "Select a Folder", "", options=options
+        )
+
+        if folder_path:
+            self.config.settings.results_path = folder_path
+            self.results_path.setText(folder_path)
+
+    def on_width_changed(self, index):
+        width_txt = self.width.itemText(index)
+        self.config.settings.lcm_diffusion_setting.image_width = int(width_txt)
+
+    def on_height_changed(self, index):
+        height_txt = self.height.itemText(index)
+        self.config.settings.lcm_diffusion_setting.image_height = int(height_txt)
+
+    def use_openvino_changed(self, state):
+        if state == 2:
+            self.lcm_model.setEnabled(False)
+            self.config.settings.lcm_diffusion_setting.use_openvino = True
+        else:
+            self.config.settings.lcm_diffusion_setting.use_openvino = False
+
+    def use_offline_model_changed(self, state):
+        if state == 2:
+            self.config.settings.lcm_diffusion_setting.use_offline_model = True
+        else:
+            self.config.settings.lcm_diffusion_setting.use_offline_model = False
+
+    def use_safety_checker_changed(self, state):
+        if state == 2:
+            self.config.settings.lcm_diffusion_setting.use_safety_checker = True
+        else:
+            self.config.settings.lcm_diffusion_setting.use_safety_checker = False
+
+    def update_steps_label(self, value):
+        self.inference_steps_value.setText(f"Number of inference steps: {value}")
+        self.config.settings.lcm_diffusion_setting.inference_steps = value
+
+    def update_num_images_label(self, value):
+        self.num_images_value.setText(f"Number of images: {value}")
+        self.config.settings.lcm_diffusion_setting.number_of_images = value
+
+    def update_guidance_label(self, value):
+        val = round(int(value) / 10, 1)
+        self.guidance_value.setText(f"Guidance scale: {val}")
+        self.config.settings.lcm_diffusion_setting.guidance_scale = val
+
+    def seed_changed(self, state):
+        if state == 2:
+            self.seed_value.setEnabled(True)
+            self.config.settings.lcm_diffusion_setting.use_seed = True
+        else:
+            self.seed_value.setEnabled(False)
+            self.config.settings.lcm_diffusion_setting.use_seed = False
+
+    def get_seed_value(self) -> int:
+        use_seed = self.config.settings.lcm_diffusion_setting.use_seed
+        seed_value = int(self.seed_value.text()) if use_seed else -1
+        return seed_value
+
+    def generate_image(self):
+        self.config.settings.lcm_diffusion_setting.seed = self.get_seed_value()
+        self.config.settings.lcm_diffusion_setting.prompt = self.prompt.toPlainText()
+
+        if self.config.settings.lcm_diffusion_setting.use_openvino:
+            model_id = LCM_DEFAULT_MODEL_OPENVINO
+        else:
+            model_id = self.lcm_model.text()
+
+        self.config.settings.lcm_diffusion_setting.lcm_model_id = model_id
+
+        reshape_required = False
+        if self.config.settings.lcm_diffusion_setting.use_openvino:
+            # Detect dimension change
+            reshape_required = is_reshape_required(
+                self.previous_width,
+                self.config.settings.lcm_diffusion_setting.image_width,
+                self.previous_height,
+                self.config.settings.lcm_diffusion_setting.image_height,
+                self.previous_model,
+                model_id,
+                self.previous_num_of_images,
+                self.config.settings.lcm_diffusion_setting.number_of_images,
+            )
+
+        images = self.context.generate_text_to_image(
+            self.config.settings,
+            reshape_required,
+            DEVICE,
+        )
+        self.image_index = 0
+        self.gen_images = []
+        for img in images:
+            im = ImageQt(img).copy()
+            pixmap = QPixmap.fromImage(im)
+            self.gen_images.append(pixmap)
+
+        if len(self.gen_images) > 1:
+            self.next_img_btn.setEnabled(True)
+            self.previous_img_btn.setEnabled(False)
+        else:
+            self.next_img_btn.setEnabled(False)
+            self.previous_img_btn.setEnabled(False)
+
+        self.img.setPixmap(self.gen_images[0])
+
+        self.previous_width = self.config.settings.lcm_diffusion_setting.image_width
+        self.previous_height = self.config.settings.lcm_diffusion_setting.image_height
+        self.previous_model = model_id
+        self.previous_num_of_images = (
+            self.config.settings.lcm_diffusion_setting.number_of_images
+        )
+
+    def text_to_image(self):
+        self.img.setText("Please wait...")
+        worker = ImageGeneratorWorker(self.generate_image)
+        self.threadpool.start(worker)
+
+    def closeEvent(self, event):
+        self.config.settings.lcm_diffusion_setting.seed = self.get_seed_value()
+        print(self.config.settings.lcm_diffusion_setting)
+        print("Saving settings")
+        self.config.save()
+
+    def reset_all_settings(self):
+        self.use_local_model_folder.setChecked(False)
+        self.width.setCurrentText("512")
+        self.height.setCurrentText("512")
+        self.inference_steps.setValue(4)
+        self.guidance.setValue(80)
+        self.use_openvino_check.setChecked(False)
+        self.seed_check.setChecked(False)
+        self.safety_checker.setChecked(True)
+        self.results_path.setText(FastStableDiffusionPaths().get_results_path())

frontend/gui/image_generator_worker.py

@@ -0,0 +1,37 @@
+from PyQt5.QtCore import (
+    pyqtSlot,
+    QRunnable,
+    pyqtSignal,
+    pyqtSlot,
+)
+from PyQt5.QtCore import QObject
+import traceback
+import sys
+
+
+class WorkerSignals(QObject):
+    finished = pyqtSignal()
+    error = pyqtSignal(tuple)
+    result = pyqtSignal(object)
+
+
+class ImageGeneratorWorker(QRunnable):
+    def __init__(self, fn, *args, **kwargs):
+        super(ImageGeneratorWorker, self).__init__()
+        self.fn = fn
+        self.args = args
+        self.kwargs = kwargs
+        self.signals = WorkerSignals()
+
+    @pyqtSlot()
+    def run(self):
+        try:
+            result = self.fn(*self.args, **self.kwargs)
+        except:
+            traceback.print_exc()
+            exctype, value = sys.exc_info()[:2]
+            self.signals.error.emit((exctype, value, traceback.format_exc()))
+        else:
+            self.signals.result.emit(result)
+        finally:
+            self.signals.finished.emit()

frontend/gui/ui.py

@@ -0,0 +1,15 @@
+from typing import List
+from frontend.gui.app_window import MainWindow
+from PyQt5.QtWidgets import QApplication
+import sys
+from app_settings import AppSettings
+
+
+def start_gui(
+    argv: List[str],
+    app_settings: AppSettings,
+):
+    app = QApplication(sys.argv)
+    window = MainWindow(app_settings)
+    window.show()
+    app.exec()

frontend/utils.py

@@ -0,0 +1,32 @@
+from constants import DEVICE
+import platform
+
+
+def is_reshape_required(
+    prev_width: int,
+    cur_width: int,
+    prev_height: int,
+    cur_height: int,
+    prev_model: int,
+    cur_model: int,
+    prev_num_of_images: int,
+    cur_num_of_images: int,
+) -> bool:
+    print(f"width - {prev_width} {cur_width}")
+    print(f"height - {prev_height} {cur_height}")
+    print(f"model - {prev_model} {cur_model}")
+    reshape_required = False
+    if (
+        prev_width != cur_width
+        or prev_height != cur_height
+        or prev_model != cur_model
+        or prev_num_of_images != cur_num_of_images
+    ):
+        print("Reshape and compile")
+        reshape_required = True
+
+    return reshape_required
+
+
+def enable_openvino_controls() -> bool:
+    return DEVICE == "cpu" and platform.system().lower() != "darwin"

frontend/webui/css/style.css

@@ -0,0 +1,24 @@
+footer {
+    visibility: hidden
+}
+
+#generate_button {
+    color: white;
+    border-color: #007bff;
+    background: #007bff;
+    width: 150px;
+    margin-top: 38px;
+    height: 80px;
+}
+
+#save_button {
+    color: white;
+    border-color: #028b40;
+    background: #01b97c;
+    width: 200px;
+}
+
+#settings_header {
+    background: rgb(245, 105, 105);
+
+}

frontend/webui/text_to_image_ui.py

@@ -0,0 +1,179 @@
+from typing import Any
+import gradio as gr
+
+from backend.models.lcmdiffusion_setting import LCMDiffusionSetting
+from context import Context
+from models.interface_types import InterfaceType
+from app_settings import Settings
+from constants import LCM_DEFAULT_MODEL, LCM_DEFAULT_MODEL_OPENVINO
+from frontend.utils import is_reshape_required
+from app_settings import AppSettings
+from constants import DEVICE
+from frontend.utils import enable_openvino_controls
+
+random_enabled = True
+
+context = Context(InterfaceType.WEBUI)
+previous_width = 0
+previous_height = 0
+previous_model_id = ""
+previous_num_of_images = 0
+
+
+def generate_text_to_image(
+    prompt,
+    image_height,
+    image_width,
+    inference_steps,
+    guidance_scale,
+    num_images,
+    seed,
+    use_openvino,
+    use_safety_checker,
+) -> Any:
+    global previous_height, previous_width, previous_model_id, previous_num_of_images
+    model_id = LCM_DEFAULT_MODEL
+    if use_openvino:
+        model_id = LCM_DEFAULT_MODEL_OPENVINO
+
+    use_seed = True if seed != -1 else False
+
+    lcm_diffusion_settings = LCMDiffusionSetting(
+        lcm_model_id=model_id,
+        prompt=prompt,
+        image_height=image_height,
+        image_width=image_width,
+        inference_steps=inference_steps,
+        guidance_scale=guidance_scale,
+        number_of_images=num_images,
+        seed=seed,
+        use_openvino=use_openvino,
+        use_safety_checker=use_safety_checker,
+        use_seed=use_seed,
+    )
+    settings = Settings(
+        lcm_diffusion_setting=lcm_diffusion_settings,
+    )
+    reshape = False
+    if use_openvino:
+        reshape = is_reshape_required(
+            previous_width,
+            image_width,
+            previous_height,
+            image_height,
+            previous_model_id,
+            model_id,
+            previous_num_of_images,
+            num_images,
+        )
+    images = context.generate_text_to_image(
+        settings,
+        reshape,
+        DEVICE,
+    )
+    previous_width = image_width
+    previous_height = image_height
+    previous_model_id = model_id
+    previous_num_of_images = num_images
+
+    return images
+
+
+def get_text_to_image_ui(app_settings: AppSettings) -> None:
+    with gr.Blocks():
+        with gr.Row():
+            with gr.Column():
+
+                def random_seed():
+                    global random_enabled
+                    random_enabled = not random_enabled
+                    seed_val = -1
+                    if not random_enabled:
+                        seed_val = 42
+                    return gr.Number.update(
+                        interactive=not random_enabled, value=seed_val
+                    )
+
+                with gr.Row():
+                    prompt = gr.Textbox(
+                        label="Describe the image you'd like to see",
+                        lines=3,
+                        placeholder="A fantasy landscape",
+                    )
+
+                    generate_btn = gr.Button(
+                        "Generate",
+                        elem_id="generate_button",
+                        scale=0,
+                    )
+                num_inference_steps = gr.Slider(
+                    1, 25, value=4, step=1, label="Inference Steps"
+                )
+                image_height = gr.Slider(
+                    256, 768, value=512, step=256, label="Image Height"
+                )
+                image_width = gr.Slider(
+                    256, 768, value=512, step=256, label="Image Width"
+                )
+                num_images = gr.Slider(
+                    1,
+                    50,
+                    value=1,
+                    step=1,
+                    label="Number of images to generate",
+                )
+                with gr.Accordion("Advanced options", open=False):
+                    guidance_scale = gr.Slider(
+                        1.0, 30.0, value=8, step=0.5, label="Guidance Scale"
+                    )
+
+                    seed = gr.Number(
+                        label="Seed",
+                        value=-1,
+                        precision=0,
+                        interactive=False,
+                    )
+                    seed_checkbox = gr.Checkbox(
+                        label="Use random seed",
+                        value=True,
+                        interactive=True,
+                    )
+
+                    openvino_checkbox = gr.Checkbox(
+                        label="Use OpenVINO",
+                        value=False,
+                        interactive=enable_openvino_controls(),
+                    )
+
+                    safety_checker_checkbox = gr.Checkbox(
+                        label="Use Safety Checker",
+                        value=True,
+                        interactive=True,
+                    )
+
+                    input_params = [
+                        prompt,
+                        image_height,
+                        image_width,
+                        num_inference_steps,
+                        guidance_scale,
+                        num_images,
+                        seed,
+                        openvino_checkbox,
+                        safety_checker_checkbox,
+                    ]
+
+            with gr.Column():
+                output = gr.Gallery(
+                    label="Generated images",
+                    show_label=True,
+                    elem_id="gallery",
+                    columns=2,
+                )
+
+    seed_checkbox.change(fn=random_seed, outputs=seed)
+    generate_btn.click(
+        fn=generate_text_to_image,
+        inputs=input_params,
+        outputs=output,
+    )

frontend/webui/ui.py

@@ -0,0 +1,36 @@
+import gradio as gr
+from constants import APP_VERSION
+from frontend.webui.text_to_image_ui import get_text_to_image_ui
+from paths import FastStableDiffusionPaths
+from app_settings import AppSettings
+
+
+def _get_footer_message() -> str:
+    version = f"<center><p> v{APP_VERSION} "
+    footer_msg = version + (
+        '  © 2023 <a href="https://github.com/rupeshs">'
+        " Rupesh Sreeraman</a></p></center>"
+    )
+    return footer_msg
+
+
+def get_web_ui(app_settings: AppSettings) -> gr.Blocks:
+    with gr.Blocks(
+        css=FastStableDiffusionPaths.get_css_path(),
+        title="FastSD CPU",
+    ) as fastsd_web_ui:
+        gr.HTML("<center><H1>FastSD CPU</H1></center>")
+        with gr.Tabs():
+            with gr.TabItem("Text to Image"):
+                get_text_to_image_ui(app_settings)
+        gr.HTML(_get_footer_message())
+
+    return fastsd_web_ui
+
+
+def start_webui(
+    app_settings: AppSettings,
+    share: bool = False,
+):
+    webui = get_web_ui(app_settings)
+    webui.launch(share=share)

models/interface_types.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class InterfaceType(Enum):
+    WEBUI = "Web User Interface"
+    GUI = "Graphical User Interface"
+    CLI = "Command Line Interface"

models/settings.py

@@ -0,0 +1,8 @@
+from pydantic import BaseModel
+from backend.models.lcmdiffusion_setting import LCMDiffusionSetting
+from paths import FastStableDiffusionPaths
+
+
+class Settings(BaseModel):
+    results_path: str = FastStableDiffusionPaths().get_results_path()
+    lcm_diffusion_setting: LCMDiffusionSetting = LCMDiffusionSetting()

paths.py

@@ -0,0 +1,48 @@
+import os
+import constants
+
+
+def join_paths(
+    first_path: str,
+    second_path: str,
+) -> str:
+    return os.path.join(first_path, second_path)
+
+
+def get_app_path():
+    app_dir = os.path.dirname(__file__)
+    work_dir = os.path.dirname(app_dir)
+    return work_dir
+
+
+def get_configs_path() -> str:
+    config_path = join_paths(get_app_path(), constants.CONFIG_DIRECTORY)
+    return config_path
+
+
+class FastStableDiffusionPaths:
+    @staticmethod
+    def get_app_settings_path() -> str:
+        configs_path = get_configs_path()
+        settings_path = join_paths(
+            configs_path,
+            constants.APP_SETTINGS_FILE,
+        )
+        return settings_path
+
+    @staticmethod
+    def get_results_path() -> str:
+        results_path = join_paths(get_app_path(), constants.RESULTS_DIRECTORY)
+        return results_path
+
+    @staticmethod
+    def get_css_path():
+        app_dir = os.path.dirname(__file__)
+        css_path = os.path.join(
+            app_dir,
+            "frontend",
+            "webui",
+            "css",
+            "style.css",
+        )
+        return css_path

requirements.txt

@@ -0,0 +1,13 @@
+accelerate==0.23.0
+diffusers==0.21.4
+transformers==4.34.0
+PyQt5
+Pillow==9.4.0
+openvino==2023.1.0
+optimum-intel==1.11.0
+onnx==1.14.1
+onnxruntime==1.16.1
+pydantic==2.4.2
+typing-extensions==4.8.0
+pyyaml
+gradio==3.39.0

utils.py

@@ -0,0 +1,10 @@
+import platform
+
+
+def show_system_info():
+    try:
+        print(f"Running on {platform.system()} platform")
+        print(f"OS: {platform.platform()}")
+        print(f"Processor: {platform.processor()}")
+    except Exception as ex:
+        print(f"Error ocurred while getting system information {ex}")