PIL.Image.Image,

np.ndarray,

torch.FloatTensor,

List[PIL.Image.Image],

List[np.ndarray],

List[torch.FloatTensor],

"""

config_name = CONFIG_NAME

@register_to_config

def __init__(

self,

do_resize: bool = True,

vae_scale_factor: int = 8,

resample: str = "lanczos",

do_normalize: bool = True,

do_binarize: bool = False,

do_convert_rgb: bool = False,

do_convert_grayscale: bool = False,

):

super().__init__()

if do_convert_rgb and do_convert_grayscale:

raise ValueError(

"`do_convert_rgb` and `do_convert_grayscale` can not both be set to `True`,"

" if you intended to convert the image into RGB format, please set `do_convert_grayscale = False`.",

" if you intended to convert the image into grayscale format, please set `do_convert_rgb = False`",

)

self.config.do_convert_rgb = False

@staticmethod

def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]:

"""

if images.ndim == 3:

images = images[None, ...]

images = (images * 255).round().astype("uint8")

if images.shape[-1] == 1:

# special case for grayscale (single channel) images

pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]

else:

pil_images = [Image.fromarray(image) for image in images]

return pil_images

@staticmethod

def pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image]) -> np.ndarray:

"""

if not isinstance(images, list):

images = [images]

images = [np.array(image).astype(np.float32) / 255.0 for image in images]

images = np.stack(images, axis=0)

return images

@staticmethod

def numpy_to_pt(images: np.ndarray) -> torch.FloatTensor:

"""

if images.ndim == 3:

images = images[..., None]

images = torch.from_numpy(images.transpose(0, 3, 1, 2))

return images

@staticmethod

def pt_to_numpy(images: torch.FloatTensor) -> np.ndarray:

"""

images = images.cpu().permute(0, 2, 3, 1).float().numpy()

return images

@staticmethod

def normalize(images: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]:

"""

return 2.0 * images - 1.0

@staticmethod

def denormalize(images: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]:

"""

return (images / 2 + 0.5).clamp(0, 1)

@staticmethod

def convert_to_rgb(image: PIL.Image.Image) -> PIL.Image.Image:

"""

image = image.convert("RGB")

return image

@staticmethod

def convert_to_grayscale(image: PIL.Image.Image) -> PIL.Image.Image:

"""

image = image.convert("L")

return image

@staticmethod

def blur(image: PIL.Image.Image, blur_factor: int = 4) -> PIL.Image.Image:

"""

image = image.filter(ImageFilter.GaussianBlur(blur_factor))

return image

@staticmethod

def get_crop_region(mask_image: PIL.Image.Image, width: int, height: int, pad=0):

"""

mask_image = mask_image.convert("L")

mask = np.array(mask_image)

# 1. find a rectangular region that contains all masked ares in an image

h, w = mask.shape

crop_left = 0

for i in range(w):

if not (mask[:, i] == 0).all():

break

crop_left += 1

crop_right = 0

for i in reversed(range(w)):

if not (mask[:, i] == 0).all():

break

crop_right += 1

crop_top = 0

for i in range(h):

if not (mask[i] == 0).all():

break

crop_top += 1

crop_bottom = 0

for i in reversed(range(h)):

if not (mask[i] == 0).all():

break

crop_bottom += 1

# 2. add padding to the crop region

x1, y1, x2, y2 = (

int(max(crop_left - pad, 0)),

int(max(crop_top - pad, 0)),

int(min(w - crop_right + pad, w)),

int(min(h - crop_bottom + pad, h)),

)

# 3. expands crop region to match the aspect ratio of the image to be processed

ratio_crop_region = (x2 - x1) / (y2 - y1)

ratio_processing = width / height

if ratio_crop_region > ratio_processing:

desired_height = (x2 - x1) / ratio_processing

desired_height_diff = int(desired_height - (y2 - y1))

y1 -= desired_height_diff // 2

y2 += desired_height_diff - desired_height_diff // 2

if y2 >= mask_image.height:

diff = y2 - mask_image.height

y2 -= diff

y1 -= diff

if y1 < 0:

y2 -= y1

y1 -= y1

if y2 >= mask_image.height:

y2 = mask_image.height

else:

desired_width = (y2 - y1) * ratio_processing

desired_width_diff = int(desired_width - (x2 - x1))

x1 -= desired_width_diff // 2

x2 += desired_width_diff - desired_width_diff // 2

if x2 >= mask_image.width:

diff = x2 - mask_image.width

x2 -= diff

x1 -= diff

if x1 < 0:

x2 -= x1

x1 -= x1

if x2 >= mask_image.width:

x2 = mask_image.width

return x1, y1, x2, y2

def _resize_and_fill(

self,

image: PIL.Image.Image,

width: int,

height: int,

) -> PIL.Image.Image:

"""

ratio = width / height

src_ratio = image.width / image.height

src_w = width if ratio < src_ratio else image.width * height // image.height

src_h = height if ratio >= src_ratio else image.height * width // image.width

resized = image.resize((src_w, src_h), resample=PIL_INTERPOLATION["lanczos"])

res = Image.new("RGB", (width, height))

res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))

if ratio < src_ratio:

fill_height = height // 2 - src_h // 2

if fill_height > 0:

res.paste(resized.resize((width, fill_height), box=(0, 0, width, 0)), box=(0, 0))

res.paste(

resized.resize((width, fill_height), box=(0, resized.height, width, resized.height)),

box=(0, fill_height + src_h),

)

elif ratio > src_ratio:

fill_width = width // 2 - src_w // 2

if fill_width > 0:

res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0))

res.paste(

resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)),

box=(fill_width + src_w, 0),

)

return res

def _resize_and_crop(

self,

image: PIL.Image.Image,

width: int,

height: int,

) -> PIL.Image.Image:

"""

ratio = width / height

src_ratio = image.width / image.height

src_w = width if ratio > src_ratio else image.width * height // image.height

src_h = height if ratio <= src_ratio else image.height * width // image.width

resized = image.resize((src_w, src_h), resample=PIL_INTERPOLATION["lanczos"])

res = Image.new("RGB", (width, height))

res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))

return res

def resize(

self,

image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],

height: int,

width: int,

resize_mode: str = "default", # "defalt", "fill", "crop"

) -> Union[PIL.Image.Image, np.ndarray, torch.Tensor]:

"""

if resize_mode != "default" and not isinstance(image, PIL.Image.Image):

raise ValueError(f"Only PIL image input is supported for resize_mode {resize_mode}")

if isinstance(image, PIL.Image.Image):

if resize_mode == "default":

image = image.resize((width, height), resample=PIL_INTERPOLATION[self.config.resample])

elif resize_mode == "fill":

image = self._resize_and_fill(image, width, height)

elif resize_mode == "crop":

image = self._resize_and_crop(image, width, height)

else:

raise ValueError(f"resize_mode {resize_mode} is not supported")

elif isinstance(image, torch.Tensor):

image = torch.nn.functional.interpolate(

image,

size=(height, width),

)

elif isinstance(image, np.ndarray):

image = self.numpy_to_pt(image)

image = torch.nn.functional.interpolate(

image,

size=(height, width),

)

image = self.pt_to_numpy(image)

return image

def binarize(self, image: PIL.Image.Image) -> PIL.Image.Image:

"""

image[image < 0.5] = 0

image[image >= 0.5] = 1

return image

def get_default_height_width(

self,

image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],

height: Optional[int] = None,

width: Optional[int] = None,

) -> Tuple[int, int]:

"""

if height is None:

if isinstance(image, PIL.Image.Image):

height = image.height

elif isinstance(image, torch.Tensor):

height = image.shape[2]

else:

height = image.shape[1]

if width is None:

if isinstance(image, PIL.Image.Image):

width = image.width

elif isinstance(image, torch.Tensor):

width = image.shape[3]

else:

width = image.shape[2]

width, height = (

x - x % self.config.vae_scale_factor for x in (width, height)

) # resize to integer multiple of vae_scale_factor

return height, width

def preprocess(

self,

image: PipelineImageInput,

height: Optional[int] = None,

width: Optional[int] = None,

resize_mode: str = "default", # "defalt", "fill", "crop"

crops_coords: Optional[Tuple[int, int, int, int]] = None,

) -> torch.Tensor:

"""

supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor)

# Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image

if self.config.do_convert_grayscale and isinstance(image, (torch.Tensor, np.ndarray)) and image.ndim == 3:

if isinstance(image, torch.Tensor):

# if image is a pytorch tensor could have 2 possible shapes:

# 1. batch x height x width: we should insert the channel dimension at position 1

# 2. channnel x height x width: we should insert batch dimension at position 0,

# however, since both channel and batch dimension has same size 1, it is same to insert at position 1

# for simplicity, we insert a dimension of size 1 at position 1 for both cases

image = image.unsqueeze(1)

else:

# if it is a numpy array, it could have 2 possible shapes:

# 1. batch x height x width: insert channel dimension on last position

# 2. height x width x channel: insert batch dimension on first position

if image.shape[-1] == 1:

image = np.expand_dims(image, axis=0)

else:

image = np.expand_dims(image, axis=-1)

if isinstance(image, supported_formats):

image = [image]

elif not (isinstance(image, list) and all(isinstance(i, supported_formats) for i in image)):

raise ValueError(

f"Input is in incorrect format: {[type(i) for i in image]}. Currently, we only support {', '.join(supported_formats)}"

)

if isinstance(image[0], PIL.Image.Image):

if crops_coords is not None:

image = [i.crop(crops_coords) for i in image]

if self.config.do_resize:

height, width = self.get_default_height_width(image[0], height, width)

image = [self.resize(i, height, width, resize_mode=resize_mode) for i in image]

if self.config.do_convert_rgb:

image = [self.convert_to_rgb(i) for i in image]

elif self.config.do_convert_grayscale:

image = [self.convert_to_grayscale(i) for i in image]

image = self.pil_to_numpy(image) # to np

image = self.numpy_to_pt(image) # to pt

elif isinstance(image[0], np.ndarray):

image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0)

image = self.numpy_to_pt(image)

height, width = self.get_default_height_width(image, height, width)

if self.config.do_resize:

image = self.resize(image, height, width)

elif isinstance(image[0], torch.Tensor):

image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0)

if self.config.do_convert_grayscale and image.ndim == 3:

image = image.unsqueeze(1)

channel = image.shape[1]

# don't need any preprocess if the image is latents

if channel == 4:

return image

height, width = self.get_default_height_width(image, height, width)

if self.config.do_resize:

image = self.resize(image, height, width)

# expected range [0,1], normalize to [-1,1]

do_normalize = self.config.do_normalize

if do_normalize and image.min() < 0:

warnings.warn(

"Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] "

f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{image.min()},{image.max()}]",

FutureWarning,

)

do_normalize = False

if do_normalize:

image = self.normalize(image)

if self.config.do_binarize:

image = self.binarize(image)

return image

def postprocess(

self,

image: torch.FloatTensor,

output_type: str = "pil",

do_denormalize: Optional[List[bool]] = None,

) -> Union[PIL.Image.Image, np.ndarray, torch.FloatTensor]:

"""

if not isinstance(image, torch.Tensor):

raise ValueError(

f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"

)

if output_type not in ["latent", "pt", "np", "pil"]:

deprecation_message = (

f"the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: "

"`pil`, `np`, `pt`, `latent`"

)

deprecate("Unsupported output_type", "1.0.0", deprecation_message, standard_warn=False)

output_type = "np"

if output_type == "latent":

return image

if do_denormalize is None:

do_denormalize = [self.config.do_normalize] * image.shape[0]

image = torch.stack(

[self.denormalize(image[i]) if do_denormalize[i] else image[i] for i in range(image.shape[0])]

)

if output_type == "pt":

return image

image = self.pt_to_numpy(image)

if output_type == "np":

return image

if output_type == "pil":

return self.numpy_to_pil(image)

def apply_overlay(

self,

mask: PIL.Image.Image,

init_image: PIL.Image.Image,

image: PIL.Image.Image,

crop_coords: Optional[Tuple[int, int, int, int]] = None,

) -> PIL.Image.Image:

"""

width, height = image.width, image.height

init_image = self.resize(init_image, width=width, height=height)

mask = self.resize(mask, width=width, height=height)

init_image_masked = PIL.Image.new("RGBa", (width, height))

init_image_masked.paste(init_image.convert("RGBA").convert("RGBa"), mask=ImageOps.invert(mask.convert("L")))

init_image_masked = init_image_masked.convert("RGBA")

if crop_coords is not None:

x, y, x2, y2 = crop_coords

w = x2 - x

h = y2 - y

base_image = PIL.Image.new("RGBA", (width, height))

image = self.resize(image, height=h, width=w, resize_mode="crop")

base_image.paste(image, (x, y))

image = base_image.convert("RGB")

image = image.convert("RGBA")

image.alpha_composite(init_image_masked)

image = image.convert("RGB")

"""

config_name = CONFIG_NAME

@register_to_config

def __init__(

self,

do_resize: bool = True,

vae_scale_factor: int = 8,

resample: str = "lanczos",

do_normalize: bool = True,

):

super().__init__()

@staticmethod

def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]:

"""

if images.ndim == 3:

images = images[None, ...]

images = (images * 255).round().astype("uint8")

if images.shape[-1] == 1:

# special case for grayscale (single channel) images

pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]

else:

pil_images = [Image.fromarray(image[:, :, :3]) for image in images]

return pil_images

@staticmethod

def depth_pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image]) -> np.ndarray:

"""

if not isinstance(images, list):

images = [images]

images = [np.array(image).astype(np.float32) / (2**16 - 1) for image in images]

images = np.stack(images, axis=0)

return images

@staticmethod

def rgblike_to_depthmap(image: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]:

"""

return image[:, :, 1] * 2**8 + image[:, :, 2]

def numpy_to_depth(self, images: np.ndarray) -> List[PIL.Image.Image]:

"""

if images.ndim == 3:

images = images[None, ...]

images_depth = images[:, :, :, 3:]

if images.shape[-1] == 6:

images_depth = (images_depth * 255).round().astype("uint8")

pil_images = [

Image.fromarray(self.rgblike_to_depthmap(image_depth), mode="I;16") for image_depth in images_depth

]

elif images.shape[-1] == 4:

images_depth = (images_depth * 65535.0).astype(np.uint16)

pil_images = [Image.fromarray(image_depth, mode="I;16") for image_depth in images_depth]

else:

raise Exception("Not supported")

return pil_images

def postprocess(

self,

image: torch.FloatTensor,

output_type: str = "pil",

do_denormalize: Optional[List[bool]] = None,

) -> Union[PIL.Image.Image, np.ndarray, torch.FloatTensor]:

"""

if not isinstance(image, torch.Tensor):

raise ValueError(

f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"

)

if output_type not in ["latent", "pt", "np", "pil"]:

deprecation_message = (

f"the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: "

"`pil`, `np`, `pt`, `latent`"

)

deprecate("Unsupported output_type", "1.0.0", deprecation_message, standard_warn=False)

output_type = "np"

if do_denormalize is None:

do_denormalize = [self.config.do_normalize] * image.shape[0]

image = torch.stack(

[self.denormalize(image[i]) if do_denormalize[i] else image[i] for i in range(image.shape[0])]

)

image = self.pt_to_numpy(image)

if output_type == "np":

if image.shape[-1] == 6:

image_depth = np.stack([self.rgblike_to_depthmap(im[:, :, 3:]) for im in image], axis=0)

else:

image_depth = image[:, :, :, 3:]

return image[:, :, :, :3], image_depth

if output_type == "pil":

return self.numpy_to_pil(image), self.numpy_to_depth(image)

else:

raise Exception(f"This type {output_type} is not supported")

def preprocess(

self,

rgb: Union[torch.FloatTensor, PIL.Image.Image, np.ndarray],

depth: Union[torch.FloatTensor, PIL.Image.Image, np.ndarray],

height: Optional[int] = None,

width: Optional[int] = None,

target_res: Optional[int] = None,

) -> torch.Tensor:

"""

supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor)

# Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image

if self.config.do_convert_grayscale and isinstance(rgb, (torch.Tensor, np.ndarray)) and rgb.ndim == 3:

raise Exception("This is not yet supported")

if isinstance(rgb, supported_formats):

rgb = [rgb]

depth = [depth]

elif not (isinstance(rgb, list) and all(isinstance(i, supported_formats) for i in rgb)):

raise ValueError(

f"Input is in incorrect format: {[type(i) for i in rgb]}. Currently, we only support {', '.join(supported_formats)}"

)

if isinstance(rgb[0], PIL.Image.Image):

if self.config.do_convert_rgb:

raise Exception("This is not yet supported")

# rgb = [self.convert_to_rgb(i) for i in rgb]

# depth = [self.convert_to_depth(i) for i in depth] #TODO define convert_to_depth

if self.config.do_resize or target_res:

height, width = self.get_default_height_width(rgb[0], height, width) if not target_res else target_res

rgb = [self.resize(i, height, width) for i in rgb]

depth = [self.resize(i, height, width) for i in depth]

rgb = self.pil_to_numpy(rgb) # to np

rgb = self.numpy_to_pt(rgb) # to pt

depth = self.depth_pil_to_numpy(depth) # to np

depth = self.numpy_to_pt(depth) # to pt

elif isinstance(rgb[0], np.ndarray):

rgb = np.concatenate(rgb, axis=0) if rgb[0].ndim == 4 else np.stack(rgb, axis=0)

rgb = self.numpy_to_pt(rgb)

height, width = self.get_default_height_width(rgb, height, width)

if self.config.do_resize:

rgb = self.resize(rgb, height, width)

depth = np.concatenate(depth, axis=0) if rgb[0].ndim == 4 else np.stack(depth, axis=0)

depth = self.numpy_to_pt(depth)

height, width = self.get_default_height_width(depth, height, width)

if self.config.do_resize:

depth = self.resize(depth, height, width)

elif isinstance(rgb[0], torch.Tensor):

raise Exception("This is not yet supported")

# rgb = torch.cat(rgb, axis=0) if rgb[0].ndim == 4 else torch.stack(rgb, axis=0)

# if self.config.do_convert_grayscale and rgb.ndim == 3:

# rgb = rgb.unsqueeze(1)

# channel = rgb.shape[1]

# height, width = self.get_default_height_width(rgb, height, width)

# if self.config.do_resize:

# rgb = self.resize(rgb, height, width)

# depth = torch.cat(depth, axis=0) if depth[0].ndim == 4 else torch.stack(depth, axis=0)

# if self.config.do_convert_grayscale and depth.ndim == 3:

# depth = depth.unsqueeze(1)

# channel = depth.shape[1]

# # don't need any preprocess if the image is latents

# if depth == 4:

# return rgb, depth

# height, width = self.get_default_height_width(depth, height, width)

# if self.config.do_resize:

# depth = self.resize(depth, height, width)

# expected range [0,1], normalize to [-1,1]

do_normalize = self.config.do_normalize

if rgb.min() < 0 and do_normalize:

warnings.warn(

"Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] "

f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{rgb.min()},{rgb.max()}]",

FutureWarning,

)

do_normalize = False

if do_normalize:

rgb = self.normalize(rgb)

depth = self.normalize(depth)

if self.config.do_binarize:

rgb = self.binarize(rgb)

depth = self.binarize(depth)