






import unittest

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

from diffusers import KandinskyPriorPipeline, 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        snake_case_				(							lowercase_   ,      unittest.TestCase   ):







   """simple docstring"""



   __lowerCAmelCase  =    KandinskyPriorPipeline
   __lowerCAmelCase  =    ["""prompt"""]
   __lowerCAmelCase  =    ["""prompt""", """negative_prompt"""]
   __lowerCAmelCase  =    [
       """num_images_per_prompt""",
       """generator""",
       """num_inference_steps""",
       """latents""",
       """negative_prompt""",
       """guidance_scale""",
       """output_type""",
       """return_dict""",
   ]
   __lowerCAmelCase  =    False

   @property
   def       SCREAMING_SNAKE_CASE__					(     self:						Any						)	->		Optional[Any]:


       """simple docstring"""



       return 32

   @property
   def       SCREAMING_SNAKE_CASE__					(     self:						Any						)	->		str:


       """simple docstring"""



       return 32

   @property
   def       SCREAMING_SNAKE_CASE__					(     self:						str						)	->		Optional[Any]:


       """simple docstring"""



       return self.time_input_dim

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


       """simple docstring"""



       return self.time_input_dim * 4

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


       """simple docstring"""



       return 100

   @property
   def       SCREAMING_SNAKE_CASE__					(     self:						Optional[Any]						)	->		Dict:


       """simple docstring"""



       _lowercase     			=						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)
       return tokenizer

   @property
   def       SCREAMING_SNAKE_CASE__					(     self:						Union[str, Any]						)	->		Optional[int]:


       """simple docstring"""



       torch.manual_seed(0						)
       _lowercase     			=						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=1_000 ,							)
       return CLIPTextModelWithProjection(snake_case__						)

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


       """simple docstring"""



       torch.manual_seed(0						)

       _lowercase     			=						{
           'num_attention_heads': 2,
           'attention_head_dim': 12,
           'embedding_dim': self.text_embedder_hidden_size,
           'num_layers': 1,
       }

       _lowercase     			=						PriorTransformer(**snake_case__						)
       # 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
       _lowercase     			=						nn.Parameter(torch.ones(model.clip_std.shape						)						)
       return model

   @property
   def       SCREAMING_SNAKE_CASE__					(     self:						Tuple						)	->		Union[str, Any]:


       """simple docstring"""



       torch.manual_seed(0						)
       _lowercase     			=						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 ,							)

       _lowercase     			=						CLIPVisionModelWithProjection(snake_case__						)
       return model

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


       """simple docstring"""



       _lowercase     			=						CLIPImageProcessor(
           crop_size=224 ,							do_center_crop=snake_case__ ,							do_normalize=snake_case__ ,							do_resize=snake_case__ ,							image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73] ,							image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11] ,							resample=3 ,							size=224 ,							)

       return image_processor

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


       """simple docstring"""



       _lowercase     			=						self.dummy_prior
       _lowercase     			=						self.dummy_image_encoder
       _lowercase     			=						self.dummy_text_encoder
       _lowercase     			=						self.dummy_tokenizer
       _lowercase     			=						self.dummy_image_processor

       _lowercase     			=						UnCLIPScheduler(
           variance_type='fixed_small_log' ,							prediction_type='sample' ,							num_train_timesteps=1_000 ,							clip_sample=snake_case__ ,							clip_sample_range=10.0 ,							)

       _lowercase     			=						{
           'prior': prior,
           'image_encoder': image_encoder,
           'text_encoder': text_encoder,
           'tokenizer': tokenizer,
           'scheduler': scheduler,
           'image_processor': image_processor,
       }

       return components

   def       SCREAMING_SNAKE_CASE__					(     self:						List[str] ,							snake_case__:						Any ,							snake_case__:						int=0						)	->		Optional[Any]:


       """simple docstring"""



       if str(snake_case__						).startswith('mps'						):
           _lowercase     			=						torch.manual_seed(snake_case__						)
       else:
           _lowercase     			=						torch.Generator(device=snake_case__						).manual_seed(snake_case__						)
       _lowercase     			=						{
           'prompt': 'horse',
           'generator': generator,
           'guidance_scale': 4.0,
           'num_inference_steps': 2,
           'output_type': 'np',
       }
       return inputs

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


       """simple docstring"""



       _lowercase     			=						'cpu'

       _lowercase     			=						self.get_dummy_components()

       _lowercase     			=						self.pipeline_class(**snake_case__						)
       _lowercase     			=						pipe.to(snake_case__						)

       pipe.set_progress_bar_config(disable=snake_case__						)

       _lowercase     			=						pipe(**self.get_dummy_inputs(snake_case__						)						)
       _lowercase     			=						output.image_embeds

       _lowercase     			=						pipe(
           **self.get_dummy_inputs(snake_case__						) ,							return_dict=snake_case__ ,							)[0]

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

       assert image.shape == (1, 32)

       _lowercase     			=						np.array(
           [-0.05_32, 1.71_20, 0.36_56, -1.08_52, -0.89_46, -1.17_56, 0.43_48, 0.24_82, 0.51_46, -0.11_56]						)

       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       SCREAMING_SNAKE_CASE__					(     self:						Any						)	->		Optional[int]:


       """simple docstring"""



       _lowercase     			=						torch_device == 'cpu'
       _lowercase     			=						True
       _lowercase     			=						False

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

   @skip_mps
   def       SCREAMING_SNAKE_CASE__					(     self:						Optional[Any]						)	->		Dict:


       """simple docstring"""



       _lowercase     			=						torch_device == 'cpu'
       _lowercase     			=						False

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








import torch

from accelerate import Accelerator, DistributedType

class        snake_case_				:







   """simple docstring"""



   def __enter__(     self:						int						)	->		Tuple:


       """simple docstring"""



       if self.enabled:
           _lowercase     			=						self.model.no_sync()
           self.model_sync_obj.__enter__()

       return self

   def __exit__(     self:						str ,							snake_case__:						Tuple ,							snake_case__:						int ,							snake_case__:						Tuple						)	->		Any:


       """simple docstring"""



       if self.enabled:
           # Average all models on exit
           self._sync_and_avg_model_params()
           self.model_sync_obj.__exit__(snake_case__ ,							snake_case__ ,							snake_case__						)

   def __init__(     self:						int ,							snake_case__:						Accelerator ,							snake_case__:						torch.nn.Module ,							snake_case__:						int ,							snake_case__:						bool = True						)	->		Optional[Any]:


       """simple docstring"""



       if accelerator.distributed_type not in [
           DistributedType.NO,
           DistributedType.MULTI_CPU,
           DistributedType.MULTI_GPU,
       ]:
           raise NotImplementedError('LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)'						)
       _lowercase     			=						enabled and accelerator.distributed_type != DistributedType.NO
       _lowercase     			=						0
       if self.enabled:
           _lowercase     			=						accelerator
           _lowercase     			=						model
           _lowercase     			=						local_sgd_steps

   def       SCREAMING_SNAKE_CASE__					(     self:						str						)	->		List[Any]:


       """simple docstring"""



       self.num_steps += 1
       if not self.enabled:
           return

       if self.num_steps % self.local_sgd_steps == 0:
           self._sync_and_avg_model_params()

   def       SCREAMING_SNAKE_CASE__					(     self:						Optional[int]						)	->		Dict:


       """simple docstring"""



       self.accelerator.wait_for_everyone()
       with self.accelerator.autocast():
           for param in self.model.parameters():
               _lowercase     			=						self.accelerator.reduce(param.data ,							reduction='mean'						)
