Upload 10 files

src/backend/base64_image.py

@@ -0,0 +1,21 @@
+from io import BytesIO
+from base64 import b64encode, b64decode
+from PIL import Image
+
+
+def pil_image_to_base64_str(
+    image: Image,
+    format: str = "JPEG",
+) -> str:
+    buffer = BytesIO()
+    image.save(buffer, format=format)
+    buffer.seek(0)
+    img_base64 = b64encode(buffer.getvalue()).decode("utf-8")
+    return img_base64
+
+
+def base64_image_to_pil(base64_str) -> Image:
+    image_data = b64decode(base64_str)
+    image_buffer = BytesIO(image_data)
+    image = Image.open(image_buffer)
+    return image

src/backend/controlnet.py

@@ -0,0 +1,90 @@
+import logging
+from PIL import Image
+from diffusers import ControlNetModel
+from backend.models.lcmdiffusion_setting import (
+    DiffusionTask,
+    ControlNetSetting,
+)
+
+
+# Prepares ControlNet adapters for use with FastSD CPU
+#
+# This function loads the ControlNet adapters defined by the
+# _lcm_diffusion_setting.controlnet_ object and returns a dictionary
+# with the pipeline arguments required to use the loaded adapters
+def load_controlnet_adapters(lcm_diffusion_setting) -> dict:
+    controlnet_args = {}
+    if (
+        lcm_diffusion_setting.controlnet is None
+        or not lcm_diffusion_setting.controlnet.enabled
+    ):
+        return controlnet_args
+
+    logging.info("Loading ControlNet adapter")
+    controlnet_adapter = ControlNetModel.from_single_file(
+        lcm_diffusion_setting.controlnet.adapter_path,
+        # local_files_only=True,
+        use_safetensors=True,
+    )
+    controlnet_args["controlnet"] = controlnet_adapter
+    return controlnet_args
+
+
+# Updates the ControlNet pipeline arguments to use for image generation
+#
+# This function uses the contents of the _lcm_diffusion_setting.controlnet_
+# object to generate a dictionary with the corresponding pipeline arguments
+# to be used for image generation; in particular, it sets the ControlNet control
+# image and conditioning scale
+def update_controlnet_arguments(lcm_diffusion_setting) -> dict:
+    controlnet_args = {}
+    if (
+        lcm_diffusion_setting.controlnet is None
+        or not lcm_diffusion_setting.controlnet.enabled
+    ):
+        return controlnet_args
+
+    controlnet_args["controlnet_conditioning_scale"] = (
+        lcm_diffusion_setting.controlnet.conditioning_scale
+    )
+    if lcm_diffusion_setting.diffusion_task == DiffusionTask.text_to_image.value:
+        controlnet_args["image"] = lcm_diffusion_setting.controlnet._control_image
+    elif lcm_diffusion_setting.diffusion_task == DiffusionTask.image_to_image.value:
+        controlnet_args["control_image"] = (
+            lcm_diffusion_setting.controlnet._control_image
+        )
+    return controlnet_args
+
+
+# Helper function to adjust ControlNet settings from a dictionary
+def controlnet_settings_from_dict(
+    lcm_diffusion_setting,
+    dictionary,
+) -> None:
+    if lcm_diffusion_setting is None or dictionary is None:
+        logging.error("Invalid arguments!")
+        return
+    if (
+        "controlnet" not in dictionary
+        or dictionary["controlnet"] is None
+        or len(dictionary["controlnet"]) == 0
+    ):
+        logging.warning("ControlNet settings not found, ControlNet will be disabled")
+        lcm_diffusion_setting.controlnet = None
+        return
+
+    controlnet = ControlNetSetting()
+    controlnet.enabled = dictionary["controlnet"][0]["enabled"]
+    controlnet.conditioning_scale = dictionary["controlnet"][0]["conditioning_scale"]
+    controlnet.adapter_path = dictionary["controlnet"][0]["adapter_path"]
+    controlnet._control_image = None
+    image_path = dictionary["controlnet"][0]["control_image"]
+    if controlnet.enabled:
+        try:
+            controlnet._control_image = Image.open(image_path)
+        except (AttributeError, FileNotFoundError) as err:
+            print(err)
+        if controlnet._control_image is None:
+            logging.error("Wrong ControlNet control image! Disabling ControlNet")
+            controlnet.enabled = False
+    lcm_diffusion_setting.controlnet = controlnet

src/backend/device.py

@@ -0,0 +1,23 @@
+import platform
+from constants import DEVICE
+import torch
+import openvino as ov
+
+core = ov.Core()
+
+
+def is_openvino_device() -> bool:
+    if DEVICE.lower() == "cpu" or DEVICE.lower()[0] == "g" or DEVICE.lower()[0] == "n":
+        return True
+    else:
+        return False
+
+
+def get_device_name() -> str:
+    if DEVICE == "cuda" or DEVICE == "mps":
+        default_gpu_index = torch.cuda.current_device()
+        return torch.cuda.get_device_name(default_gpu_index)
+    elif platform.system().lower() == "darwin":
+        return platform.processor()
+    elif is_openvino_device():
+        return core.get_property(DEVICE.upper(), "FULL_DEVICE_NAME")

src/backend/image_saver.py

@@ -0,0 +1,75 @@
+import json
+from os import path, mkdir
+from typing import Any
+from uuid import uuid4
+from backend.models.lcmdiffusion_setting import LCMDiffusionSetting
+from utils import get_image_file_extension
+
+
+def get_exclude_keys():
+    exclude_keys = {
+        "init_image": True,
+        "generated_images": True,
+        "lora": {
+            "models_dir": True,
+            "path": True,
+        },
+        "dirs": True,
+        "controlnet": {
+            "adapter_path": True,
+        },
+    }
+    return exclude_keys
+
+
+class ImageSaver:
+    @staticmethod
+    def save_images(
+        output_path: str,
+        images: Any,
+        folder_name: str = "",
+        format: str = "PNG",
+        jpeg_quality: int = 90,
+        lcm_diffusion_setting: LCMDiffusionSetting = None,
+    ) -> list[str]:
+        gen_id = uuid4()
+        image_ids = []
+
+        if images:
+            image_seeds = []
+
+            for index, image in enumerate(images):
+
+                image_seed = image.info.get('image_seed')
+                if image_seed is not None:
+                    image_seeds.append(image_seed)
+
+                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_extension = get_image_file_extension(format)
+                image_file_name = f"{gen_id}-{index+1}{image_extension}"
+                image_ids.append(image_file_name)
+                image.save(path.join(out_path, image_file_name), quality = jpeg_quality)
+            if lcm_diffusion_setting:
+                data = lcm_diffusion_setting.model_dump(exclude=get_exclude_keys())
+                if image_seeds:
+                    data['image_seeds'] = image_seeds
+                with open(path.join(out_path, f"{gen_id}.json"), "w") as json_file:
+                    json.dump(
+                        data,
+                        json_file,
+                        indent=4,
+                    )
+        return image_ids
+            

src/backend/lcm_text_to_image.py

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

src/backend/lora.py

@@ -0,0 +1,136 @@
+import glob
+from os import path
+from paths import get_file_name, FastStableDiffusionPaths
+from pathlib import Path
+
+
+# A basic class to keep track of the currently loaded LoRAs and
+# their weights; the diffusers function \c get_active_adapters()
+# returns a list of adapter names but not their weights so we need
+# a way to keep track of the current LoRA weights to set whenever
+# a new LoRA is loaded
+class _lora_info:
+    def __init__(
+        self,
+        path: str,
+        weight: float,
+    ):
+        self.path = path
+        self.adapter_name = get_file_name(path)
+        self.weight = weight
+
+    def __del__(self):
+        self.path = None
+        self.adapter_name = None
+
+
+_loaded_loras = []
+_current_pipeline = None
+
+
+# This function loads a LoRA from the LoRA path setting, so it's
+# possible to load multiple LoRAs by calling this function more than
+# once with a different LoRA path setting; note that if you plan to
+# load multiple LoRAs and dynamically change their weights, you
+# might want to set the LoRA fuse option to False
+def load_lora_weight(
+    pipeline,
+    lcm_diffusion_setting,
+):
+    if not lcm_diffusion_setting.lora.path:
+        raise Exception("Empty lora model path")
+
+    if not path.exists(lcm_diffusion_setting.lora.path):
+        raise Exception("Lora model path is invalid")
+
+    # If the pipeline has been rebuilt since the last call, remove all
+    # references to previously loaded LoRAs and store the new pipeline
+    global _loaded_loras
+    global _current_pipeline
+    if pipeline != _current_pipeline:
+        for lora in _loaded_loras:
+            del lora
+        del _loaded_loras
+        _loaded_loras = []
+        _current_pipeline = pipeline
+
+    current_lora = _lora_info(
+        lcm_diffusion_setting.lora.path,
+        lcm_diffusion_setting.lora.weight,
+    )
+    _loaded_loras.append(current_lora)
+
+    if lcm_diffusion_setting.lora.enabled:
+        print(f"LoRA adapter name : {current_lora.adapter_name}")
+        pipeline.load_lora_weights(
+            FastStableDiffusionPaths.get_lora_models_path(),
+            weight_name=Path(lcm_diffusion_setting.lora.path).name,
+            local_files_only=True,
+            adapter_name=current_lora.adapter_name,
+        )
+        update_lora_weights(
+            pipeline,
+            lcm_diffusion_setting,
+        )
+
+        if lcm_diffusion_setting.lora.fuse:
+            pipeline.fuse_lora()
+
+
+def get_lora_models(root_dir: str):
+    lora_models = glob.glob(f"{root_dir}/**/*.safetensors", recursive=True)
+    lora_models_map = {}
+    for file_path in lora_models:
+        lora_name = get_file_name(file_path)
+        if lora_name is not None:
+            lora_models_map[lora_name] = file_path
+    return lora_models_map
+
+
+# This function returns a list of (adapter_name, weight) tuples for the
+# currently loaded LoRAs
+def get_active_lora_weights():
+    active_loras = []
+    for lora_info in _loaded_loras:
+        active_loras.append(
+            (
+                lora_info.adapter_name,
+                lora_info.weight,
+            )
+        )
+    return active_loras
+
+
+# This function receives a pipeline, an lcm_diffusion_setting object and
+# an optional list of updated (adapter_name, weight) tuples
+def update_lora_weights(
+    pipeline,
+    lcm_diffusion_setting,
+    lora_weights=None,
+):
+    global _loaded_loras
+    global _current_pipeline
+    if pipeline != _current_pipeline:
+        print("Wrong pipeline when trying to update LoRA weights")
+        return
+    if lora_weights:
+        for idx, lora in enumerate(lora_weights):
+            if _loaded_loras[idx].adapter_name != lora[0]:
+                print("Wrong adapter name in LoRA enumeration!")
+                continue
+            _loaded_loras[idx].weight = lora[1]
+
+    adapter_names = []
+    adapter_weights = []
+    if lcm_diffusion_setting.use_lcm_lora:
+        adapter_names.append("lcm")
+        adapter_weights.append(1.0)
+    for lora in _loaded_loras:
+        adapter_names.append(lora.adapter_name)
+        adapter_weights.append(lora.weight)
+    pipeline.set_adapters(
+        adapter_names,
+        adapter_weights=adapter_weights,
+    )
+    adapter_weights = zip(adapter_names, adapter_weights)
+    print(f"Adapters: {list(adapter_weights)}")

src/backend/safety_checker.py

@@ -0,0 +1,29 @@
+from typing import Any
+
+from transformers import pipeline
+
+from constants import SAFETY_CHECKER_MODEL
+
+
+class SafetyChecker:
+    """A class to check if an image is NSFW or not."""
+
+    def __init__(
+        self,
+        mode_id: str = SAFETY_CHECKER_MODEL,
+    ):
+        self.classifier = pipeline(
+            "image-classification",
+            model=mode_id,
+        )
+
+    def is_safe(
+        self,
+        image: Any,
+    ) -> bool:
+        pred = self.classifier(image)
+        scores = {label["label"]: label["score"] for label in pred}
+        nsfw_score = scores.get("nsfw", 0)
+        normal_score = scores.get("normal", 0)
+        print(f"NSFW score: {nsfw_score}, Normal score: {normal_score}")
+        return normal_score > nsfw_score

src/backend/tiny_autoencoder.py

@@ -0,0 +1,40 @@
+from constants import (
+    TAESD_MODEL,
+    TAESDXL_MODEL,
+    TAESD_MODEL_OPENVINO,
+    TAESDXL_MODEL_OPENVINO,
+    TAEF1_MODEL_OPENVINO,
+)
+
+
+def get_tiny_autoencoder_repo_id(pipeline_class) -> str:
+    print(f"Pipeline class : {pipeline_class}")
+    if (
+        pipeline_class == "LatentConsistencyModelPipeline"
+        or pipeline_class == "StableDiffusionPipeline"
+        or pipeline_class == "StableDiffusionImg2ImgPipeline"
+        or pipeline_class == "StableDiffusionControlNetPipeline"
+        or pipeline_class == "StableDiffusionControlNetImg2ImgPipeline"
+    ):
+        return TAESD_MODEL
+    elif (
+        pipeline_class == "StableDiffusionXLPipeline"
+        or pipeline_class == "StableDiffusionXLImg2ImgPipeline"
+    ):
+        return TAESDXL_MODEL
+    elif (
+        pipeline_class == "OVStableDiffusionPipeline"
+        or pipeline_class == "OVStableDiffusionImg2ImgPipeline"
+    ):
+        return TAESD_MODEL_OPENVINO
+    elif (
+        pipeline_class == "OVStableDiffusionXLPipeline"
+        or pipeline_class == "OVStableDiffusionXLImg2ImgPipeline"
+    ):
+        return TAESDXL_MODEL_OPENVINO
+    elif pipeline_class == "OVFluxPipeline":
+        return TAEF1_MODEL_OPENVINO
+    else:
+        raise ValueError(
+            f"Tiny autoencoder not available for the pipeline class {pipeline_class}!"
+        )

src/backend/utils.py

@@ -0,0 +1,18 @@
+from PIL import Image
+
+
+def get_blank_image(
+    width: int,
+    height: int,
+) -> Image.Image:
+    """
+    Create a blank image with the specified width and height.
+
+    Args:
+        width (int): The width of the image.
+        height (int): The height of the image.
+
+    Returns:
+        Image.Image: A blank image with the specified dimensions.
+    """
+    return Image.new("RGB", (width, height), (0, 0, 0))