| from pprint import pprint |
| from time import perf_counter |
| from traceback import print_exc |
| from typing import Any |
|
|
| from app_settings import Settings |
| from backend.image_saver import ImageSaver |
| from backend.lcm_text_to_image import LCMTextToImage |
| from backend.models.lcmdiffusion_setting import DiffusionTask |
| from backend.utils import get_blank_image |
| from models.interface_types import InterfaceType |
|
|
|
|
| class Context: |
| def __init__( |
| self, |
| interface_type: InterfaceType, |
| device="cpu", |
| ): |
| self.interface_type = interface_type.value |
| self.lcm_text_to_image = LCMTextToImage(device) |
| self._latency = 0 |
| self._error = "" |
|
|
| @property |
| def latency(self): |
| return self._latency |
|
|
| @property |
| def error(self): |
| return self._error |
|
|
| def generate_text_to_image( |
| self, |
| settings: Settings, |
| reshape: bool = False, |
| device: str = "cpu", |
| save_config=True, |
| ) -> Any: |
| try: |
| self._error = "" |
| tick = perf_counter() |
| from state import get_settings |
|
|
| if ( |
| settings.lcm_diffusion_setting.diffusion_task |
| == DiffusionTask.text_to_image.value |
| ): |
| settings.lcm_diffusion_setting.init_image = None |
|
|
| if save_config: |
| get_settings().save() |
|
|
| pprint(settings.lcm_diffusion_setting.model_dump()) |
| if not settings.lcm_diffusion_setting.lcm_lora: |
| return None |
| self.lcm_text_to_image.init( |
| device, |
| settings.lcm_diffusion_setting, |
| ) |
|
|
| images = self.lcm_text_to_image.generate( |
| settings.lcm_diffusion_setting, |
| reshape, |
| ) |
|
|
| elapsed = perf_counter() - tick |
| self._latency = elapsed |
| print(f"Latency : {elapsed:.2f} seconds") |
| if settings.lcm_diffusion_setting.controlnet: |
| if settings.lcm_diffusion_setting.controlnet.enabled: |
| images.append( |
| settings.lcm_diffusion_setting.controlnet._control_image |
| ) |
|
|
| if settings.lcm_diffusion_setting.use_safety_checker: |
| print("Safety Checker is enabled") |
| from state import get_safety_checker |
|
|
| safety_checker = get_safety_checker() |
| blank_image = get_blank_image( |
| settings.lcm_diffusion_setting.image_width, |
| settings.lcm_diffusion_setting.image_height, |
| ) |
| for idx, image in enumerate(images): |
| if not safety_checker.is_safe(image): |
| images[idx] = blank_image |
| except Exception as exception: |
| print(f"Error in generating images: {exception}") |
| self._error = str(exception) |
| print_exc() |
| return None |
| return images |
|
|
| def save_images( |
| self, |
| images: Any, |
| settings: Settings, |
| ) -> list[str]: |
| saved_images = [] |
| if images and settings.generated_images.save_image: |
| saved_images = ImageSaver.save_images( |
| settings.generated_images.path, |
| images=images, |
| lcm_diffusion_setting=settings.lcm_diffusion_setting, |
| format=settings.generated_images.format, |
| jpeg_quality=settings.generated_images.save_image_quality, |
| ) |
| return saved_images |
|
|
