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