






import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class 				_A			(     a__					,					unittest.TestCase				):
    """simple docstring"""

    lowerCamelCase	:							Optional[Any]				= KandinskyVaaPriorPipeline
    lowerCamelCase	:							List[str]				= ["prompt"]
    lowerCamelCase	:							Union[str, Any]				= ["prompt", "negative_prompt"]
    lowerCamelCase	:							Optional[int]				= [
        "num_images_per_prompt",
        "generator",
        "num_inference_steps",
        "latents",
        "negative_prompt",
        "guidance_scale",
        "output_type",
        "return_dict",
    ]
    lowerCamelCase	:							str				= False



    @property
    def _a							(				self      :     Union[str, Any]      )      ->							Dict:
           return 32



    @property
    def _a							(				self      :     Optional[Any]      )      ->							str:
           return 32



    @property
    def _a							(				self      :     int      )      ->							int:
           return self.time_input_dim



    @property
    def _a							(				self      :     Tuple      )      ->							Union[str, Any]:
           return self.time_input_dim * 4



    @property
    def _a							(				self      :     List[Any]      )      ->							List[Any]:
           return 100



    @property
    def _a							(				self      :     Optional[Any]      )      ->							Union[str, Any]:
           __UpperCAmelCase     				=CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""      )
           return tokenizer



    @property
    def _a							(				self      :     Optional[int]      )      ->							List[str]:
           torch.manual_seed(0      )
           __UpperCAmelCase     				=CLIPTextConfig(
               bos_token_id=0 ,     eos_token_id=2 ,     hidden_size=self.text_embedder_hidden_size ,     projection_dim=self.text_embedder_hidden_size ,     intermediate_size=37 ,     layer_norm_eps=1e-05 ,     num_attention_heads=4 ,     num_hidden_layers=5 ,     pad_token_id=1 ,     vocab_size=1000 ,     )
           return CLIPTextModelWithProjection(lowerCamelCase_      )



    @property
    def _a							(				self      :     Optional[int]      )      ->							List[str]:
           torch.manual_seed(0      )

           __UpperCAmelCase     				={
               """num_attention_heads""": 2,
               """attention_head_dim""": 12,
               """embedding_dim""": self.text_embedder_hidden_size,
               """num_layers""": 1,
           }

           __UpperCAmelCase     				=PriorTransformer(**lowerCamelCase_      )
           # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
           __UpperCAmelCase     				=nn.Parameter(torch.ones(model.clip_std.shape      )      )
           return model



    @property
    def _a							(				self      :     Tuple      )      ->							int:
           torch.manual_seed(0      )
           __UpperCAmelCase     				=CLIPVisionConfig(
               hidden_size=self.text_embedder_hidden_size ,     image_size=224 ,     projection_dim=self.text_embedder_hidden_size ,     intermediate_size=37 ,     num_attention_heads=4 ,     num_channels=3 ,     num_hidden_layers=5 ,     patch_size=14 ,     )

           __UpperCAmelCase     				=CLIPVisionModelWithProjection(lowerCamelCase_      )
           return model



    @property
    def _a							(				self      :     Optional[Any]      )      ->							Any:
           __UpperCAmelCase     				=CLIPImageProcessor(
               crop_size=224 ,     do_center_crop=lowerCamelCase_ ,     do_normalize=lowerCamelCase_ ,     do_resize=lowerCamelCase_ ,     image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073] ,     image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711] ,     resample=3 ,     size=224 ,     )

           return image_processor



    def _a							(				self      :     Optional[int]      )      ->							Dict:
           __UpperCAmelCase     				=self.dummy_prior
           __UpperCAmelCase     				=self.dummy_image_encoder
           __UpperCAmelCase     				=self.dummy_text_encoder
           __UpperCAmelCase     				=self.dummy_tokenizer
           __UpperCAmelCase     				=self.dummy_image_processor

           __UpperCAmelCase     				=UnCLIPScheduler(
               variance_type="""fixed_small_log""" ,     prediction_type="""sample""" ,     num_train_timesteps=1000 ,     clip_sample=lowerCamelCase_ ,     clip_sample_range=10.0 ,     )

           __UpperCAmelCase     				={
               """prior""": prior,
               """image_encoder""": image_encoder,
               """text_encoder""": text_encoder,
               """tokenizer""": tokenizer,
               """scheduler""": scheduler,
               """image_processor""": image_processor,
           }

           return components



    def _a							(				self      :     Optional[int] ,     __SCREAMING_SNAKE_CASE      :     Union[str, Any] ,     __SCREAMING_SNAKE_CASE      :     str=0      )      ->							int:
           if str(lowerCamelCase_      ).startswith("""mps"""      ):
                  __UpperCAmelCase     				=torch.manual_seed(lowerCamelCase_      )
           else:
                  __UpperCAmelCase     				=torch.Generator(device=lowerCamelCase_      ).manual_seed(lowerCamelCase_      )
           __UpperCAmelCase     				={
               """prompt""": """horse""",
               """generator""": generator,
               """guidance_scale""": 4.0,
               """num_inference_steps""": 2,
               """output_type""": """np""",
           }
           return inputs



    def _a							(				self      :     List[Any]      )      ->							str:
           __UpperCAmelCase     				="""cpu"""

           __UpperCAmelCase     				=self.get_dummy_components()

           __UpperCAmelCase     				=self.pipeline_class(**lowerCamelCase_      )
           __UpperCAmelCase     				=pipe.to(lowerCamelCase_      )

           pipe.set_progress_bar_config(disable=lowerCamelCase_      )

           __UpperCAmelCase     				=pipe(**self.get_dummy_inputs(lowerCamelCase_      )      )
           __UpperCAmelCase     				=output.image_embeds

           __UpperCAmelCase     				=pipe(
               **self.get_dummy_inputs(lowerCamelCase_      ) ,     return_dict=lowerCamelCase_ ,     )[0]

           __UpperCAmelCase     				=image[0, -10:]
           __UpperCAmelCase     				=image_from_tuple[0, -10:]

           assert image.shape == (1, 32)

           __UpperCAmelCase     				=np.array(
               [-0.0_532, 1.7_120, 0.3_656, -1.0_852, -0.8_946, -1.1_756, 0.4_348, 0.2_482, 0.5_146, -0.1_156]      )

           assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2
           assert np.abs(image_from_tuple_slice.flatten() - expected_slice      ).max() < 1e-2



    @skip_mps
    def _a							(				self      :     Optional[int]      )      ->							Tuple:
           __UpperCAmelCase     				=torch_device == """cpu"""
           __UpperCAmelCase     				=True
           __UpperCAmelCase     				=False

           self._test_inference_batch_single_identical(
               test_max_difference=lowerCamelCase_ ,     relax_max_difference=lowerCamelCase_ ,     test_mean_pixel_difference=lowerCamelCase_ ,     )



    @skip_mps
    def _a							(				self      :     Any      )      ->							Tuple:
           __UpperCAmelCase     				=torch_device == """cpu"""
           __UpperCAmelCase     				=False

           self._test_attention_slicing_forward_pass(
               test_max_difference=lowerCamelCase_ ,     test_mean_pixel_difference=lowerCamelCase_ ,     )








'''simple docstring'''






import unittest

import torch
from torch import nn

from accelerate.test_utils import require_cuda
from accelerate.utils.memory import find_executable_batch_size, release_memory


def    _snake_case  (	)		->       Union[str, Any]:
							raise RuntimeError('''CUDA out of memory.'''			)




class 						a__							(				nn.Module    ):





							'''simple docstring'''






							def __init__(						self  )			->				int:
														super().__init__()
														lowerCAmelCase__			  =      nn.Linear(3   ,       4  )
														lowerCAmelCase__			  =      nn.BatchNormad(4  )
														lowerCAmelCase__			  =      nn.Linear(4   ,       5  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_  )			->				Optional[Any]:
														return self.lineara(self.batchnorm(self.lineara(lowerCamelCase_  )  )  )




class 						a__							(				unittest.TestCase    ):





							'''simple docstring'''






							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														lowerCAmelCase__			  =      []

														@find_executable_batch_size(starting_batch_size=1_28  )
														def mock_training_loop_function(lowerCamelCase_  ):
																					nonlocal batch_sizes
																					batch_sizes.append(lowerCamelCase_  )
																					if batch_size != 8:
																												raise_fake_out_of_memory()

														mock_training_loop_function()
														self.assertListEqual(lowerCamelCase_   ,       [1_28, 64, 32, 16, 8]  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														lowerCAmelCase__			  =      []

														@find_executable_batch_size(starting_batch_size=1_28  )
														def mock_training_loop_function(lowerCamelCase_   ,       lowerCamelCase_  ):
																					nonlocal batch_sizes
																					batch_sizes.append(lowerCamelCase_  )
																					if batch_size != 8:
																												raise_fake_out_of_memory()
																					return batch_size, arga

														lowerCAmelCase__  ,   lowerCAmelCase__			  =      mock_training_loop_function('''hello'''  )
														self.assertListEqual(lowerCamelCase_   ,       [1_28, 64, 32, 16, 8]  )
														self.assertListEqual([bs, arga]   ,       [8, '''hello''']  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														@find_executable_batch_size(starting_batch_size=0  )
														def mock_training_loop_function(lowerCamelCase_  ):
																					pass

														with self.assertRaises(lowerCamelCase_  ) as cm:
																					mock_training_loop_function()
																					self.assertIn('''No executable batch size found, reached zero.'''   ,       cm.exception.args[0]  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Dict:
														@find_executable_batch_size(starting_batch_size=16  )
														def mock_training_loop_function(lowerCamelCase_  ):
																					if batch_size > 0:
																												raise_fake_out_of_memory()
																					pass

														with self.assertRaises(lowerCamelCase_  ) as cm:
																					mock_training_loop_function()
																					self.assertIn('''No executable batch size found, reached zero.'''   ,       cm.exception.args[0]  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														@find_executable_batch_size(starting_batch_size=1_28  )
														def mock_training_loop_function(lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  ):
																					if batch_size != 8:
																												raise raise_fake_out_of_memory()

														with self.assertRaises(lowerCamelCase_  ) as cm:
																					mock_training_loop_function(1_28   ,       '''hello'''   ,       '''world'''  )
																					self.assertIn('''Batch size was passed into `f`'''   ,       cm.exception.args[0]  )
																					self.assertIn('''`f(arg1=\'hello\', arg2=\'world\')'''   ,       cm.exception.args[0]  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Union[str, Any]:
														@find_executable_batch_size(starting_batch_size=16  )
														def mock_training_loop_function(lowerCamelCase_  ):
																					raise ValueError('''Oops, we had an error!'''  )

														with self.assertRaises(lowerCamelCase_  ) as cm:
																					mock_training_loop_function()
																					self.assertIn('''Oops, we had an error!'''   ,       cm.exception.args[0]  )







							@require_cuda
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[Any]:
														lowerCAmelCase__			  =      torch.cuda.memory_allocated()
														lowerCAmelCase__			  =      ModelForTest()
														model.cuda()
														self.assertGreater(torch.cuda.memory_allocated()   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      release_memory(lowerCamelCase_  )
														self.assertEqual(torch.cuda.memory_allocated()   ,       lowerCamelCase_  )