|
|
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 |
|
|
|
