

"""simple docstring"""


from typing import Callable, List, Optional, Union

import PIL
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPSegForImageSegmentation,
    CLIPSegProcessor,
    CLIPTextModel,
    CLIPTokenizer,
)

from diffusers import DiffusionPipeline
from diffusers.configuration_utils import FrozenDict
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionInpaintPipeline
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import deprecate, is_accelerate_available, logging


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





class    snake_case     (    __snake_case  ):

				def __init__(     self							:    Optional[Any]					,  UpperCamelCase__							:    CLIPSegForImageSegmentation					,  UpperCamelCase__							:    CLIPSegProcessor					,  UpperCamelCase__							:    AutoencoderKL					,  UpperCamelCase__							:    CLIPTextModel					,  UpperCamelCase__							:    CLIPTokenizer					,  UpperCamelCase__							:    UNetaDConditionModel					,  UpperCamelCase__							:    Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]					,  UpperCamelCase__							:    StableDiffusionSafetyChecker					,  UpperCamelCase__							:    CLIPImageProcessor					,  )->       List[Any]:



						'''simple docstring'''


						super().__init__()

						if hasattr(scheduler.config					,  "steps_offset") and scheduler.config.steps_offset != 1:
								__lowerCAmelCase:	Optional[Any]   				=      (
								    f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
								    f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
								    "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
								    " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
								    " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
								    " file"
								)
								deprecate("steps_offset!=1"					,  "1.0.0"					,  lowercase_					,  standard_warn=lowercase_)
								__lowerCAmelCase:	List[Any]   				=      dict(scheduler.config)
								__lowerCAmelCase:	Optional[Any]   				=      1
								__lowerCAmelCase:	Optional[Any]   				=      FrozenDict(lowercase_)

						if hasattr(scheduler.config					,  "skip_prk_steps") and scheduler.config.skip_prk_steps is False:
								__lowerCAmelCase:	str   				=      (
								    f"The configuration file of this scheduler: {scheduler} has not set the configuration"
								    " `skip_prk_steps`. `skip_prk_steps` should be set to True in the configuration file. Please make"
								    " sure to update the config accordingly as not setting `skip_prk_steps` in the config might lead to"
								    " incorrect results in future versions. If you have downloaded this checkpoint from the Hugging Face"
								    " Hub, it would be very nice if you could open a Pull request for the"
								    " `scheduler/scheduler_config.json` file"
								)
								deprecate("skip_prk_steps not set"					,  "1.0.0"					,  lowercase_					,  standard_warn=lowercase_)
								__lowerCAmelCase:	List[str]   				=      dict(scheduler.config)
								__lowerCAmelCase:	Any   				=      True
								__lowerCAmelCase:	Optional[Any]   				=      FrozenDict(lowercase_)

						if safety_checker is None:
								logger.warning(
								    f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
								    " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
								    " results in services or applications open to the public. Both the diffusers team and Hugging Face"
								    " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
								    " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
								    " information, please have a look at https://github.com/huggingface/diffusers/pull/254 .")

						self.register_modules(
						    segmentation_model=lowercase_					,  segmentation_processor=lowercase_					,  vae=lowercase_					,  text_encoder=lowercase_					,  tokenizer=lowercase_					,  unet=lowercase_					,  scheduler=lowercase_					,  safety_checker=lowercase_					,  feature_extractor=lowercase_					,  )

				def 							lowercase_					(     self							:    int					,  UpperCamelCase__							:    Optional[Union[str, int]] = "auto")->       Optional[int]:



						'''simple docstring'''


						if slice_size == "auto":
								# half the attention head size is usually a good trade-off between
								# speed and memory
								__lowerCAmelCase:	str   				=      self.unet.config.attention_head_dim // 2
						self.unet.set_attention_slice(lowercase_)

				def 							lowercase_					(     self							:    Tuple)->       Optional[Any]:



						'''simple docstring'''


						self.enable_attention_slicing(lowercase_)

				def 							lowercase_					(     self							:    Dict)->       str:



						'''simple docstring'''


						if is_accelerate_available():
								from accelerate import cpu_offload
						else:
								raise ImportError("Please install accelerate via `pip install accelerate`")

						__lowerCAmelCase:	Tuple   				=      torch.device("cuda")

						for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.safety_checker]:
								if cpu_offloaded_model is not None:
										cpu_offload(lowercase_					,  lowercase_)

				@property
				# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
				def 							lowercase_					(     self							:    Tuple)->       int:



						'''simple docstring'''


						if self.device != torch.device("meta") or not hasattr(self.unet					,  "_hf_hook"):
								return self.device
						for module in self.unet.modules():
								if (
								    hasattr(lowercase_					,  "_hf_hook")
								    and hasattr(module._hf_hook					,  "execution_device")
								    and module._hf_hook.execution_device is not None
								):
										return torch.device(module._hf_hook.execution_device)
						return self.device




				@torch.no_grad()
				def __call__(     self							:    int					,  UpperCamelCase__							:    Union[str, List[str]]					,  UpperCamelCase__							:    Union[torch.FloatTensor, PIL.Image.Image]					,  UpperCamelCase__							:    str					,  UpperCamelCase__							:    int = 5_1_2					,  UpperCamelCase__							:    int = 5_1_2					,  UpperCamelCase__							:    int = 5_0					,  UpperCamelCase__							:    float = 7.5					,  UpperCamelCase__							:    Optional[Union[str, List[str]]] = None					,  UpperCamelCase__							:    Optional[int] = 1					,  UpperCamelCase__							:    float = 0.0					,  UpperCamelCase__							:    Optional[torch.Generator] = None					,  UpperCamelCase__							:    Optional[torch.FloatTensor] = None					,  UpperCamelCase__							:    Optional[str] = "pil"					,  UpperCamelCase__							:    bool = True					,  UpperCamelCase__							:    Optional[Callable[[int, int, torch.FloatTensor], None]] = None					,  UpperCamelCase__							:    int = 1					,  **UpperCamelCase__							:    List[Any]					,  )->       Tuple:



						'''simple docstring'''


						__lowerCAmelCase:	List[str]   				=      self.segmentation_processor(
						    text=[text]					,  images=[image]					,  padding="max_length"					,  return_tensors="pt").to(self.device)
						__lowerCAmelCase:	str   				=      self.segmentation_model(**lowercase_)
						__lowerCAmelCase:	str   				=      torch.sigmoid(outputs.logits).cpu().detach().unsqueeze(-1).numpy()
						__lowerCAmelCase:	int   				=      self.numpy_to_pil(lowercase_)[0].resize(image.size)

						# Run inpainting pipeline with the generated mask
						__lowerCAmelCase:	Optional[int]   				=      StableDiffusionInpaintPipeline(
						    vae=self.vae					,  text_encoder=self.text_encoder					,  tokenizer=self.tokenizer					,  unet=self.unet					,  scheduler=self.scheduler					,  safety_checker=self.safety_checker					,  feature_extractor=self.feature_extractor					,  )
						return inpainting_pipeline(
						    prompt=lowercase_					,  image=lowercase_					,  mask_image=lowercase_					,  height=lowercase_					,  width=lowercase_					,  num_inference_steps=lowercase_					,  guidance_scale=lowercase_					,  negative_prompt=lowercase_					,  num_images_per_prompt=lowercase_					,  eta=lowercase_					,  generator=lowercase_					,  latents=lowercase_					,  output_type=lowercase_					,  return_dict=lowercase_					,  callback=lowercase_					,  callback_steps=lowercase_					,  )





"""simple docstring"""


import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()





class    snake_case     (    __snake_case,							unittest.TestCase  ):
				SCREAMING_SNAKE_CASE_       :				Tuple       				=  PriorTransformer
				SCREAMING_SNAKE_CASE_       :				List[str]       				=  """hidden_states"""

				@property
				def 							lowercase_					(     self							:    Dict)->       str:



						'''simple docstring'''


						__lowerCAmelCase:	str   				=      4
						__lowerCAmelCase:	int   				=      8
						__lowerCAmelCase:	int   				=      7

						__lowerCAmelCase:	str   				=      floats_tensor((batch_size, embedding_dim)).to(UpperCamelCase__)

						__lowerCAmelCase:	Optional[Any]   				=      floats_tensor((batch_size, embedding_dim)).to(UpperCamelCase__)
						__lowerCAmelCase:	Any   				=      floats_tensor((batch_size, num_embeddings, embedding_dim)).to(UpperCamelCase__)

						return {
						    "hidden_states": hidden_states,
						    "timestep": 2,
						    "proj_embedding": proj_embedding,
						    "encoder_hidden_states": encoder_hidden_states,
						}

				def 							lowercase_					(     self							:    Optional[int]					,  UpperCamelCase__							:    str=0)->       str:



						'''simple docstring'''


						torch.manual_seed(UpperCamelCase__)
						__lowerCAmelCase:	List[Any]   				=      4
						__lowerCAmelCase:	Dict   				=      8
						__lowerCAmelCase:	int   				=      7

						__lowerCAmelCase:	List[str]   				=      torch.randn((batch_size, embedding_dim)).to(UpperCamelCase__)

						__lowerCAmelCase:	Tuple   				=      torch.randn((batch_size, embedding_dim)).to(UpperCamelCase__)
						__lowerCAmelCase:	List[Any]   				=      torch.randn((batch_size, num_embeddings, embedding_dim)).to(UpperCamelCase__)

						return {
						    "hidden_states": hidden_states,
						    "timestep": 2,
						    "proj_embedding": proj_embedding,
						    "encoder_hidden_states": encoder_hidden_states,
						}

				@property
				def 							lowercase_					(     self							:    Dict)->       List[Any]:



						'''simple docstring'''


						return (4, 8)

				@property
				def 							lowercase_					(     self							:    Optional[int])->       int:



						'''simple docstring'''


						return (4, 8)

				def 							lowercase_					(     self							:    Optional[int])->       Tuple:



						'''simple docstring'''


						__lowerCAmelCase:	str   				=      {
						    "num_attention_heads": 2,
						    "attention_head_dim": 4,
						    "num_layers": 2,
						    "embedding_dim": 8,
						    "num_embeddings": 7,
						    "additional_embeddings": 4,
						}
						__lowerCAmelCase:	Any   				=      self.dummy_input
						return init_dict, inputs_dict

				def 							lowercase_					(     self							:    List[Any])->       int:



						'''simple docstring'''


						__lowerCAmelCase				,			__lowerCAmelCase:	Optional[int]   				=      PriorTransformer.from_pretrained(
						    "hf-internal-testing/prior-dummy"					,  output_loading_info=UpperCamelCase__)
						self.assertIsNotNone(UpperCamelCase__)
						self.assertEqual(len(loading_info["missing_keys"])					,  0)

						model.to(UpperCamelCase__)
						__lowerCAmelCase:	Dict   				=      model(**self.dummy_input)[0]

						assert hidden_states is not None, "Make sure output is not None"

				def 							lowercase_					(     self							:    List[str])->       Tuple:



						'''simple docstring'''


						__lowerCAmelCase				,			__lowerCAmelCase:	Optional[Any]   				=      self.prepare_init_args_and_inputs_for_common()

						__lowerCAmelCase:	Tuple   				=      self.model_class(**UpperCamelCase__)
						__lowerCAmelCase:	List[str]   				=      inspect.signature(model.forward)
						# signature.parameters is an OrderedDict => so arg_names order is deterministic
						__lowerCAmelCase:	List[Any]   				=      [*signature.parameters.keys()]

						__lowerCAmelCase:	Any   				=      ["hidden_states", "timestep"]
						self.assertListEqual(arg_names[:2]					,  UpperCamelCase__)




				def 							lowercase_					(     self							:    Optional[int])->       List[str]:



						'''simple docstring'''


						__lowerCAmelCase:	int   				=      PriorTransformer.from_pretrained("hf-internal-testing/prior-dummy")
						__lowerCAmelCase:	Union[str, Any]   				=      model.to(UpperCamelCase__)

						if hasattr(UpperCamelCase__					,  "set_default_attn_processor"):
								model.set_default_attn_processor()

						__lowerCAmelCase:	str   				=      self.get_dummy_seed_input()

						with torch.no_grad():
								__lowerCAmelCase:	Dict   				=      model(**UpperCamelCase__)[0]

						__lowerCAmelCase:	Dict   				=      output[0, :5].flatten().cpu()
						print(UpperCamelCase__)

						# Since the VAE Gaussian prior's generator is seeded on the appropriate device,
						# the expected output slices are not the same for CPU and GPU.
						__lowerCAmelCase:	List[str]   				=      torch.tensor([-1.3436, -0.2870, 0.7538, 0.4368, -0.0239])
						self.assertTrue(torch_all_close(UpperCamelCase__					,  UpperCamelCase__					,  rtol=1e-2))





@slow
class    snake_case     (    unittest.TestCase  ):

				def 							lowercase_					(     self							:    int					,  UpperCamelCase__							:    Dict=1					,  UpperCamelCase__							:    str=7_6_8					,  UpperCamelCase__							:    int=7_7					,  UpperCamelCase__							:    Any=0)->       Union[str, Any]:



						'''simple docstring'''


						torch.manual_seed(UpperCamelCase__)
						__lowerCAmelCase:	List[Any]   				=      batch_size
						__lowerCAmelCase:	Any   				=      embedding_dim
						__lowerCAmelCase:	Dict   				=      num_embeddings

						__lowerCAmelCase:	Dict   				=      torch.randn((batch_size, embedding_dim)).to(UpperCamelCase__)

						__lowerCAmelCase:	str   				=      torch.randn((batch_size, embedding_dim)).to(UpperCamelCase__)
						__lowerCAmelCase:	int   				=      torch.randn((batch_size, num_embeddings, embedding_dim)).to(UpperCamelCase__)

						return {
						    "hidden_states": hidden_states,
						    "timestep": 2,
						    "proj_embedding": proj_embedding,
						    "encoder_hidden_states": encoder_hidden_states,
						}

				def 							lowercase_					(     self							:    List[Any])->       Union[str, Any]:



						'''simple docstring'''


						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()




				@parameterized.expand(
				    [
				        # fmt: off
				        [1_3, [-0.5861,  0.1283, -0.0931,  0.0882,  0.4476,  0.1329, -0.0498,  0.0640]],
				        [3_7, [-0.4913,  0.0110, -0.0483,  0.0541,  0.4954, -0.0170,  0.0354,  0.1651]],
				        # fmt: on
				    ])
				def 							lowercase_					(     self							:    Optional[int]					,  UpperCamelCase__							:    List[str]					,  UpperCamelCase__							:    int)->       List[Any]:



						'''simple docstring'''


						__lowerCAmelCase:	List[str]   				=      PriorTransformer.from_pretrained("kandinsky-community/kandinsky-2-1-prior"					,  subfolder="prior")
						model.to(UpperCamelCase__)
						__lowerCAmelCase:	Dict   				=      self.get_dummy_seed_input(seed=UpperCamelCase__)

						with torch.no_grad():
								__lowerCAmelCase:	Optional[Any]   				=      model(**UpperCamelCase__)[0]

						assert list(sample.shape) == [1, 7_6_8]

						__lowerCAmelCase:	Dict   				=      sample[0, :8].flatten().cpu()
						print(UpperCamelCase__)
						__lowerCAmelCase:	Union[str, Any]   				=      torch.tensor(UpperCamelCase__)

						assert torch_all_close(UpperCamelCase__					,  UpperCamelCase__					,  atol=1e-3)


