Spaces:
Paused
Paused
File size: 3,405 Bytes
99b955f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
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
|