


"""simple docstring"""






from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


_lowerCamelCase     =  logging.get_logger(__name__)


if is_vision_available():
	import PIL


class 			lowerCamelCase_		(			lowercase   ):





				"""simple docstring"""







				_lowerCAmelCase      :      Optional[Any]      							=					["pixel_values"]



				def __init__(						self       ,		UpperCAmelCase__ = True       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = PILImageResampling.BICUBIC       ,		UpperCAmelCase__ = True       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = True       ,		UpperCAmelCase__ = 1 / 255       ,		UpperCAmelCase__ = True       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = True       ,		**UpperCAmelCase__       ,		):
								super().__init__(**UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				size if size is not None else {"shortest_edge": 224}
								SCREAMING_SNAKE_CASE__     							=				get_size_dict(UpperCAmelCase__       ,		default_to_square=UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				crop_size if crop_size is not None else {"height": 224, "width": 224}
								SCREAMING_SNAKE_CASE__     							=				get_size_dict(UpperCAmelCase__       ,		default_to_square=UpperCAmelCase__       ,		param_name="crop_size"     )

								SCREAMING_SNAKE_CASE__     							=				do_resize
								SCREAMING_SNAKE_CASE__     							=				size
								SCREAMING_SNAKE_CASE__     							=				resample
								SCREAMING_SNAKE_CASE__     							=				do_center_crop
								SCREAMING_SNAKE_CASE__     							=				crop_size
								SCREAMING_SNAKE_CASE__     							=				do_rescale
								SCREAMING_SNAKE_CASE__     							=				rescale_factor
								SCREAMING_SNAKE_CASE__     							=				do_normalize
								SCREAMING_SNAKE_CASE__     							=				image_mean if image_mean is not None else OPENAI_CLIP_MEAN
								SCREAMING_SNAKE_CASE__     							=				image_std if image_std is not None else OPENAI_CLIP_STD
								SCREAMING_SNAKE_CASE__     							=				do_convert_rgb



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__ = PILImageResampling.BICUBIC       ,		UpperCAmelCase__ = None       ,		**UpperCAmelCase__       ,		):
								SCREAMING_SNAKE_CASE__     							=				get_size_dict(UpperCAmelCase__       ,		default_to_square=UpperCAmelCase__     )
								if "shortest_edge" not in size:
												raise ValueError(f'''The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}'''     )
								SCREAMING_SNAKE_CASE__     							=				get_resize_output_image_size(UpperCAmelCase__       ,		size=size["shortest_edge"]       ,		default_to_square=UpperCAmelCase__     )
								return resize(UpperCAmelCase__       ,		size=UpperCAmelCase__       ,		resample=UpperCAmelCase__       ,		data_format=UpperCAmelCase__       ,		**UpperCAmelCase__     )



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__ = None       ,		**UpperCAmelCase__       ,		):
								SCREAMING_SNAKE_CASE__     							=				get_size_dict(UpperCAmelCase__     )
								if "height" not in size or "width" not in size:
												raise ValueError(f'''The `size` parameter must contain the keys (height, width). Got {size.keys()}'''     )
								return center_crop(UpperCAmelCase__       ,		size=(size["height"], size["width"])       ,		data_format=UpperCAmelCase__       ,		**UpperCAmelCase__     )



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__ = None       ,		**UpperCAmelCase__       ,		):
								return rescale(UpperCAmelCase__       ,		scale=UpperCAmelCase__       ,		data_format=UpperCAmelCase__       ,		**UpperCAmelCase__     )



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__ = None       ,		**UpperCAmelCase__       ,		):
								return normalize(UpperCAmelCase__       ,		mean=UpperCAmelCase__       ,		std=UpperCAmelCase__       ,		data_format=UpperCAmelCase__       ,		**UpperCAmelCase__     )






				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = None       ,		UpperCAmelCase__ = ChannelDimension.FIRST       ,		**UpperCAmelCase__       ,		):
								SCREAMING_SNAKE_CASE__     							=				do_resize if do_resize is not None else self.do_resize
								SCREAMING_SNAKE_CASE__     							=				size if size is not None else self.size
								SCREAMING_SNAKE_CASE__     							=				get_size_dict(UpperCAmelCase__       ,		param_name="size"       ,		default_to_square=UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				resample if resample is not None else self.resample
								SCREAMING_SNAKE_CASE__     							=				do_center_crop if do_center_crop is not None else self.do_center_crop
								SCREAMING_SNAKE_CASE__     							=				crop_size if crop_size is not None else self.crop_size
								SCREAMING_SNAKE_CASE__     							=				get_size_dict(UpperCAmelCase__       ,		param_name="crop_size"       ,		default_to_square=UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				do_rescale if do_rescale is not None else self.do_rescale
								SCREAMING_SNAKE_CASE__     							=				rescale_factor if rescale_factor is not None else self.rescale_factor
								SCREAMING_SNAKE_CASE__     							=				do_normalize if do_normalize is not None else self.do_normalize
								SCREAMING_SNAKE_CASE__     							=				image_mean if image_mean is not None else self.image_mean
								SCREAMING_SNAKE_CASE__     							=				image_std if image_std is not None else self.image_std
								SCREAMING_SNAKE_CASE__     							=				do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

								SCREAMING_SNAKE_CASE__     							=				make_list_of_images(UpperCAmelCase__     )

								if not valid_images(UpperCAmelCase__     ):
												raise ValueError(
												    "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
												    "torch.Tensor, tf.Tensor or jax.ndarray."     )

								if do_resize and size is None:
												raise ValueError("Size must be specified if do_resize is True."     )

								if do_center_crop and crop_size is None:
												raise ValueError("Crop size must be specified if do_center_crop is True."     )

								if do_rescale and rescale_factor is None:
												raise ValueError("Rescale factor must be specified if do_rescale is True."     )

								if do_normalize and (image_mean is None or image_std is None):
												raise ValueError("Image mean and std must be specified if do_normalize is True."     )

								# PIL RGBA images are converted to RGB
								if do_convert_rgb:
												SCREAMING_SNAKE_CASE__     							=				[convert_to_rgb(UpperCAmelCase__     ) for image in images]

								# All transformations expect numpy arrays.
								SCREAMING_SNAKE_CASE__     							=				[to_numpy_array(UpperCAmelCase__     ) for image in images]

								if do_resize:
												SCREAMING_SNAKE_CASE__     							=				[self.resize(image=UpperCAmelCase__       ,		size=UpperCAmelCase__       ,		resample=UpperCAmelCase__     ) for image in images]

								if do_center_crop:
												SCREAMING_SNAKE_CASE__     							=				[self.center_crop(image=UpperCAmelCase__       ,		size=UpperCAmelCase__     ) for image in images]

								if do_rescale:
												SCREAMING_SNAKE_CASE__     							=				[self.rescale(image=UpperCAmelCase__       ,		scale=UpperCAmelCase__     ) for image in images]

								if do_normalize:
												SCREAMING_SNAKE_CASE__     							=				[self.normalize(image=UpperCAmelCase__       ,		mean=UpperCAmelCase__       ,		std=UpperCAmelCase__     ) for image in images]

								SCREAMING_SNAKE_CASE__     							=				[to_channel_dimension_format(UpperCAmelCase__       ,		UpperCAmelCase__     ) for image in images]

								SCREAMING_SNAKE_CASE__     							=				{"pixel_values": images}
								return BatchFeature(data=UpperCAmelCase__       ,		tensor_type=UpperCAmelCase__     )




"""simple docstring"""






import numpy as np
import torch
from imwatermark import WatermarkEncoder


# Copied from https://github.com/Stability-AI/generative-models/blob/613af104c6b85184091d42d374fef420eddb356d/scripts/demo/streamlit_helpers.py#L66
_lowerCamelCase     =  0b1011_0011_1110_1100_1001_0000_0111_1011_1011_0001_1001_1110
# bin(x)[2:] gives bits of x as str, use int to convert them to 0/1
_lowerCamelCase     =  [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]


class 			lowerCamelCase_		:





				"""simple docstring"""







				def __init__(						self     ):
								SCREAMING_SNAKE_CASE__     							=				WATERMARK_BITS
								SCREAMING_SNAKE_CASE__     							=				WatermarkEncoder()

								self.encoder.set_watermark("bits"       ,		self.watermark     )






				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__     ):
								# can't encode images that are smaller than 256
								if images.shape[-1] < 256:
												return images

								SCREAMING_SNAKE_CASE__     							=				(255 * (images / 2 + 0.5)).cpu().permute(0       ,		2       ,		3       ,		1     ).float().numpy()

								SCREAMING_SNAKE_CASE__     							=				[self.encoder.encode(UpperCAmelCase__       ,		"dwtDct"     ) for image in images]

								SCREAMING_SNAKE_CASE__     							=				torch.from_numpy(np.array(UpperCAmelCase__     )     ).permute(0       ,		3       ,		1       ,		2     )

								SCREAMING_SNAKE_CASE__     							=				torch.clamp(2 * (images / 255 - 0.5)       ,		min=-1.0       ,		max=1.0     )
								return images
