


from typing import List, Optional, Tuple, Union

from . import logging
from .import_utils import is_torch_available, is_torch_version


if is_torch_available():
			import torch

__UpperCamelCase:  Tuple   =   logging.get_logger(__name__)  # pylint: disable=invalid-name

try:
			from torch._dynamo import allow_in_graph as maybe_allow_in_graph
except (ImportError, ModuleNotFoundError):


			def 			SCREAMING_SNAKE_CASE__     (					cls    :   Optional[int] )      ->       Union[str, Any]:




					'''simple docstring'''
					return cls


def 			SCREAMING_SNAKE_CASE__     (					_lowercase    :   Union[Tuple, List]      ,						_lowercase    :   Optional[Union[List["torch.Generator"], "torch.Generator"]] = None      ,						_lowercase    :   Optional["torch.device"] = None      ,						_lowercase    :   Optional["torch.dtype"] = None      ,						_lowercase    :   Optional["torch.layout"] = None      ,						)      ->       Optional[Any]:




		'''simple docstring'''
		lowercase__				:			List[Any]            =		device
		lowercase__				:			List[str]            =		shape[0]

		lowercase__				:			str            =		layout or torch.strided
		lowercase__				:			Tuple            =		device or torch.device('cpu' )

		if generator is not None:
				lowercase__				:			Union[str, Any]            =		generator.device.type if not isinstance(_lowercase      ,						_lowercase ) else generator[0].device.type
				if gen_device_type != device.type and gen_device_type == "cpu":
						lowercase__				:			Union[str, Any]            =		'cpu'
						if device != "mps":
								logger.info(
								    f"""The passed generator was created on 'cpu' even though a tensor on {device} was expected."""
								    f""" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably"""
								    f""" slighly speed up this function by passing a generator that was created on the {device} device.""" )
				elif gen_device_type != device.type and gen_device_type == "cuda":
						raise ValueError(f"""Cannot generate a {device} tensor from a generator of type {gen_device_type}.""" )

		if isinstance(_lowercase      ,						_lowercase ):
				lowercase__				:			List[str]            =		(1,) + shape[1:]
				lowercase__				:			str            =		[
				    torch.randn(_lowercase      ,						generator=generator[i]      ,						device=_lowercase      ,						dtype=_lowercase      ,						layout=_lowercase )
				    for i in range(_lowercase )
				]
				lowercase__				:			List[str]            =		torch.cat(_lowercase      ,						dim=0 ).to(_lowercase )
		else:
				lowercase__				:			Union[str, Any]            =		torch.randn(_lowercase      ,						generator=_lowercase      ,						device=_lowercase      ,						dtype=_lowercase      ,						layout=_lowercase ).to(_lowercase )

		return latents





def 			SCREAMING_SNAKE_CASE__     (					_lowercase    :   int )      ->       Union[str, Any]:




		'''simple docstring'''
		if is_torch_version('<'      ,						'2.0.0' ) or not hasattr(_lowercase      ,						'_dynamo' ):
				return False
		return isinstance(_lowercase      ,						torch._dynamo.eval_frame.OptimizedModule )









from typing import List

import PIL
import torch
from PIL import Image

from ...configuration_utils import ConfigMixin
from ...models.modeling_utils import ModelMixin
from ...utils import PIL_INTERPOLATION


class 			__lowerCAmelCase				(					_UpperCamelCase      ,						_UpperCamelCase			):


	'''simple docstring'''





	def __init__(	self:	Union[str, Any] ):
			super().__init__()

			self.register_buffer('watermark_image',				torch.zeros((62, 62, 4) ) )
			lowercase__				:			int            =		None





	def       snake_case__(	self:	Optional[int],				lowerCamelCase_:	List[PIL.Image.Image],				lowerCamelCase_:	Dict=None ):
			# copied from https://github.com/deep-floyd/IF/blob/b77482e36ca2031cb94dbca1001fc1e6400bf4ab/deepfloyd_if/modules/base.py#L287

			lowercase__				:			str            =		images[0].height
			lowercase__				:			Optional[Any]            =		images[0].width

			lowercase__				:			List[Any]            =		sample_size or h

			lowercase__				:			Any            =		min(h / sample_size,				w / sample_size )
			lowercase__      ,							lowercase__				:			Optional[int]            =		(int(h / coef ), int(w / coef )) if coef < 1 else (h, w)

			lowercase__      ,							lowercase__				:			Any            =		1024**2, img_w * img_h
			lowercase__				:			Union[str, Any]            =		(Sa / Sa) ** 0.5
			lowercase__      ,							lowercase__      ,							lowercase__				:			Optional[int]            =		int(K * 62 ), img_w - int(14 * K ), img_h - int(14 * K )

			if self.watermark_image_as_pil is None:
					lowercase__				:			Tuple            =		self.watermark_image.to(torch.uinta ).cpu().numpy()
					lowercase__				:			str            =		Image.fromarray(lowerCamelCase_,				mode='RGBA' )
					lowercase__				:			Tuple            =		watermark_image

			lowercase__				:			Optional[int]            =		self.watermark_image_as_pil.resize(
			    (wm_size, wm_size),				PIL_INTERPOLATION['bicubic'],				reducing_gap=lowerCamelCase_ )

			for pil_img in images:
					pil_img.paste(lowerCamelCase_,				box=(wm_x - wm_size, wm_y - wm_size, wm_x, wm_y),				mask=wm_img.split()[-1] )

			return images


