




"""simple docstring"""





import random
import unittest

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

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

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()

class       __lowercase       (			snake_case_						,   unittest.TestCase	):


				'''simple docstring'''


				_lowerCAmelCase				=KandinskyVaaPriorEmbaEmbPipeline
				_lowerCAmelCase				=["prompt", "image"]
				_lowerCAmelCase				=["prompt", "image"]
				_lowerCAmelCase				=[
				    "num_images_per_prompt",
				    "strength",
				    "generator",
				    "num_inference_steps",
				    "latents",
				    "negative_prompt",
				    "guidance_scale",
				    "output_type",
				    "return_dict",
				]
				_lowerCAmelCase				=False

				@property
				def       _lowerCamelCase					(   self   :     List[str]      ):
										return 3_2

				@property
				def       _lowerCamelCase					(   self   :     List[str]      ):
										return 3_2

				@property
				def       _lowerCamelCase					(   self   :     int      ):
										return self.time_input_dim

				@property
				def       _lowerCamelCase					(   self   :     int      ):
										return self.time_input_dim * 4

				@property
				def       _lowerCamelCase					(   self   :     Optional[int]      ):
										return 1_0_0

				@property
				def       _lowerCamelCase					(   self   :     Union[str, Any]      ):
										snake_case__       :					List[str]										=					CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      )
										return tokenizer

				@property
				def       _lowerCamelCase					(   self   :     Dict      ):
										torch.manual_seed(0      )
										snake_case__       :					Optional[int]										=					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=3_7      ,layer_norm_eps=1E-05      ,num_attention_heads=4      ,num_hidden_layers=5      ,pad_token_id=1      ,vocab_size=1_0_0_0      ,)
										return CLIPTextModelWithProjection(SCREAMING_SNAKE_CASE_      )

				@property
				def       _lowerCamelCase					(   self   :     Union[str, Any]      ):
										torch.manual_seed(0      )

										snake_case__       :					Dict										=					{
										    '''num_attention_heads''': 2,
										    '''attention_head_dim''': 1_2,
										    '''embedding_dim''': self.text_embedder_hidden_size,
										    '''num_layers''': 1,
										}

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

				@property
				def       _lowerCamelCase					(   self   :     Tuple      ):
										torch.manual_seed(0      )
										snake_case__       :					Dict										=					CLIPVisionConfig(
										    hidden_size=self.text_embedder_hidden_size      ,image_size=2_2_4      ,projection_dim=self.text_embedder_hidden_size      ,intermediate_size=3_7      ,num_attention_heads=4      ,num_channels=3      ,num_hidden_layers=5      ,patch_size=1_4      ,)

										snake_case__       :					Union[str, Any]										=					CLIPVisionModelWithProjection(SCREAMING_SNAKE_CASE_      )
										return model

				@property
				def       _lowerCamelCase					(   self   :     Optional[int]      ):
										snake_case__       :					Tuple										=					CLIPImageProcessor(
										    crop_size=2_2_4      ,do_center_crop=SCREAMING_SNAKE_CASE_      ,do_normalize=SCREAMING_SNAKE_CASE_      ,do_resize=SCREAMING_SNAKE_CASE_      ,image_mean=[0.48145466, 0.4578275, 0.40821073]      ,image_std=[0.26862954, 0.26130258, 0.27577711]      ,resample=3      ,size=2_2_4      ,)

										return image_processor

				def       _lowerCamelCase					(   self   :     Dict      ):
										snake_case__       :					Optional[Any]										=					self.dummy_prior
										snake_case__       :					int										=					self.dummy_image_encoder
										snake_case__       :					int										=					self.dummy_text_encoder
										snake_case__       :					List[str]										=					self.dummy_tokenizer
										snake_case__       :					List[Any]										=					self.dummy_image_processor

										snake_case__       :					Dict										=					UnCLIPScheduler(
										    variance_type='''fixed_small_log'''      ,prediction_type='''sample'''      ,num_train_timesteps=1_0_0_0      ,clip_sample=SCREAMING_SNAKE_CASE_      ,clip_sample_range=10.0      ,)

										snake_case__       :					Union[str, Any]										=					{
										    '''prior''': prior,
										    '''image_encoder''': image_encoder,
										    '''text_encoder''': text_encoder,
										    '''tokenizer''': tokenizer,
										    '''scheduler''': scheduler,
										    '''image_processor''': image_processor,
										}

										return components

				def       _lowerCamelCase					(   self   :     Optional[int]      ,SCREAMING_SNAKE_CASE_   :     str      ,SCREAMING_SNAKE_CASE_   :     Optional[Any]=0      ):
										if str(SCREAMING_SNAKE_CASE_      ).startswith('''mps'''      ):
																snake_case__       :					List[Any]										=					torch.manual_seed(SCREAMING_SNAKE_CASE_      )
										else:
																snake_case__       :					Union[str, Any]										=					torch.Generator(device=SCREAMING_SNAKE_CASE_      ).manual_seed(SCREAMING_SNAKE_CASE_      )

										snake_case__       :					Dict										=					floats_tensor((1, 3, 6_4, 6_4)      ,rng=random.Random(SCREAMING_SNAKE_CASE_      )      ).to(SCREAMING_SNAKE_CASE_      )
										snake_case__       :					Tuple										=					image.cpu().permute(0      ,2      ,3      ,1      )[0]
										snake_case__       :					Any										=					Image.fromarray(np.uinta(SCREAMING_SNAKE_CASE_      )      ).convert('''RGB'''      ).resize((2_5_6, 2_5_6)      )

										snake_case__       :					Tuple										=					{
										    '''prompt''': '''horse''',
										    '''image''': init_image,
										    '''strength''': 0.5,
										    '''generator''': generator,
										    '''guidance_scale''': 4.0,
										    '''num_inference_steps''': 2,
										    '''output_type''': '''np''',
										}
										return inputs

				def       _lowerCamelCase					(   self   :     Optional[Any]      ):
										snake_case__       :					Tuple										=					'''cpu'''

										snake_case__       :					List[str]										=					self.get_dummy_components()

										snake_case__       :					Union[str, Any]										=					self.pipeline_class(**SCREAMING_SNAKE_CASE_      )
										snake_case__       :					Tuple										=					pipe.to(SCREAMING_SNAKE_CASE_      )

										pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_      )

										snake_case__       :					str										=					pipe(**self.get_dummy_inputs(SCREAMING_SNAKE_CASE_      )      )
										snake_case__       :					Any										=					output.image_embeds

										snake_case__       :					str										=					pipe(
										    **self.get_dummy_inputs(SCREAMING_SNAKE_CASE_      )      ,return_dict=SCREAMING_SNAKE_CASE_      ,)[0]

										snake_case__       :					str										=					image[0, -1_0:]
										snake_case__       :					str										=					image_from_tuple[0, -1_0:]

										assert image.shape == (1, 3_2)

										snake_case__       :					str										=					np.array(
										    [
										        0.1071284,
										        1.3330271,
										        0.61260223,
										        -0.6691065,
										        -0.3846852,
										        -1.0303661,
										        0.22716111,
										        0.03348901,
										        0.30040675,
										        -0.24805029,
										    ]      )

										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       _lowerCamelCase					(   self   :     List[Any]      ):
										snake_case__       :					List[str]										=					torch_device == '''cpu'''
										snake_case__       :					int										=					True
										snake_case__       :					Optional[Any]										=					False

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





				@skip_mps
				def       _lowerCamelCase					(   self   :     List[str]      ):
										snake_case__       :					Optional[Any]										=					torch_device == '''cpu'''
										snake_case__       :					List[Any]										=					False

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








"""simple docstring"""





import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
							import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


_lowercase	      =    get_logger(__name__)

def 		lowerCAmelCase							( lowerCamelCase      ,			lowerCamelCase      ,			lowerCamelCase      ,			lowerCamelCase      ,			lowerCamelCase      ,			lowerCamelCase = None      ,			):
						'''simple docstring'''


						for i, state in enumerate(lowerCamelCase		):
												snake_case__       :					List[Any]										=					f'''{MODEL_NAME}.bin''' if i == 0 else f'''{MODEL_NAME}_{i}.bin'''
												snake_case__       :					Tuple										=					os.path.join(lowerCamelCase      ,			lowerCamelCase		)
												save(lowerCamelCase      ,			lowerCamelCase		)
												logger.info(f'''Model weights saved in {output_model_file}'''		)
						# Optimizer states
						for i, opt in enumerate(lowerCamelCase		):
												snake_case__       :					str										=					opt.state_dict()
												snake_case__       :					Optional[Any]										=					f'''{OPTIMIZER_NAME}.bin''' if i == 0 else f'''{OPTIMIZER_NAME}_{i}.bin'''
												snake_case__       :					str										=					os.path.join(lowerCamelCase      ,			lowerCamelCase		)
												save(lowerCamelCase      ,			lowerCamelCase		)
												logger.info(f'''Optimizer state saved in {output_optimizer_file}'''		)
						# Scheduler states
						for i, scheduler in enumerate(lowerCamelCase		):
												snake_case__       :					int										=					scheduler.state_dict()
												snake_case__       :					List[Any]										=					f'''{SCHEDULER_NAME}.bin''' if i == 0 else f'''{SCHEDULER_NAME}_{i}.bin'''
												snake_case__       :					int										=					os.path.join(lowerCamelCase      ,			lowerCamelCase		)
												save(lowerCamelCase      ,			lowerCamelCase		)
												logger.info(f'''Scheduler state saved in {output_scheduler_file}'''		)
						# GradScaler state
						if scaler is not None:
												snake_case__       :					Tuple										=					scaler.state_dict()
												snake_case__       :					Tuple										=					os.path.join(lowerCamelCase      ,			lowerCamelCase		)
												torch.save(lowerCamelCase      ,			lowerCamelCase		)
												logger.info(f'''Gradient scaler state saved in {output_scaler_file}'''		)
						# Random number generator states
						snake_case__       :					Dict										=					{}
						snake_case__       :					Dict										=					f'''{RNG_STATE_NAME}_{process_index}.pkl'''
						snake_case__       :					Dict										=					random.getstate()
						snake_case__       :					List[str]										=					np.random.get_state()
						snake_case__       :					str										=					torch.get_rng_state()
						if is_xpu_available():
												snake_case__       :					List[Any]										=					torch.xpu.get_rng_state_all()
						else:
												snake_case__       :					Union[str, Any]										=					torch.cuda.get_rng_state_all()
						if is_tpu_available():
												snake_case__       :					int										=					xm.get_rng_state()
						snake_case__       :					int										=					os.path.join(lowerCamelCase      ,			lowerCamelCase		)
						torch.save(lowerCamelCase      ,			lowerCamelCase		)
						logger.info(f'''Random states saved in {output_states_file}'''		)
						return output_dir

def 		lowerCAmelCase							( lowerCamelCase      ,			lowerCamelCase      ,			lowerCamelCase      ,			lowerCamelCase      ,			lowerCamelCase      ,			lowerCamelCase=None      ,			lowerCamelCase=None      ,			**lowerCamelCase      ,			):
						'''simple docstring'''


						if map_location not in [None, "cpu", "on_device"]:
												raise TypeError(
												    '''Unsupported optimizer map location passed, please choose one of `None`, `\'cpu\'`, or `\'on_device\'`'''		)
						if map_location is None:
												snake_case__       :					str										=					'''cpu'''
						elif map_location == "on_device":
												snake_case__       :					List[Any]										=					PartialState().device
						# Model states
						for i, model in enumerate(lowerCamelCase		):
												snake_case__       :					Optional[int]										=					f'''{MODEL_NAME}.bin''' if i == 0 else f'''{MODEL_NAME}_{i}.bin'''
												snake_case__       :					int										=					os.path.join(lowerCamelCase      ,			lowerCamelCase		)
												models[i].load_state_dict(torch.load(lowerCamelCase      ,			map_location=lowerCamelCase		)      ,			**lowerCamelCase		)
						logger.info('''All model weights loaded successfully'''		)

						# Optimizer states
						for i, opt in enumerate(lowerCamelCase		):
												snake_case__       :					Any										=					f'''{OPTIMIZER_NAME}.bin''' if i == 0 else f'''{OPTIMIZER_NAME}_{i}.bin'''
												snake_case__       :					Optional[Any]										=					os.path.join(lowerCamelCase      ,			lowerCamelCase		)
												snake_case__       :					Any										=					torch.load(lowerCamelCase      ,			map_location=lowerCamelCase		)
												optimizers[i].load_state_dict(lowerCamelCase		)
						logger.info('''All optimizer states loaded successfully'''		)

						# Scheduler states
						for i, scheduler in enumerate(lowerCamelCase		):
												snake_case__       :					Optional[int]										=					f'''{SCHEDULER_NAME}.bin''' if i == 0 else f'''{SCHEDULER_NAME}_{i}.bin'''
												snake_case__       :					Optional[Any]										=					os.path.join(lowerCamelCase      ,			lowerCamelCase		)
												scheduler.load_state_dict(torch.load(lowerCamelCase		)		)
						logger.info('''All scheduler states loaded successfully'''		)

						# GradScaler state
						if scaler is not None:
												snake_case__       :					Optional[Any]										=					os.path.join(lowerCamelCase      ,			lowerCamelCase		)
												scaler.load_state_dict(torch.load(lowerCamelCase		)		)
												logger.info('''GradScaler state loaded successfully'''		)

						# Random states
						try:
												snake_case__       :					Optional[int]										=					torch.load(os.path.join(lowerCamelCase      ,			f'''{RNG_STATE_NAME}_{process_index}.pkl'''		)		)
												random.setstate(states['''random_state''']		)
												np.random.set_state(states['''numpy_random_seed''']		)
												torch.set_rng_state(states['''torch_manual_seed''']		)
												if is_xpu_available():
																		torch.xpu.set_rng_state_all(states['''torch_xpu_manual_seed''']		)
												else:
																		torch.cuda.set_rng_state_all(states['''torch_cuda_manual_seed''']		)
												if is_tpu_available():
																		xm.set_rng_state(states['''xm_seed''']		)
												logger.info('''All random states loaded successfully'''		)
						except Exception:
												logger.info('''Could not load random states'''		)

def 		lowerCAmelCase							( lowerCamelCase      ,			lowerCamelCase      ,			lowerCamelCase = 0		):
						'''simple docstring'''


						snake_case__       :					Union[str, Any]										=					Path(lowerCamelCase		) / f'''custom_checkpoint_{index}.pkl'''
						logger.info(f'''Saving the state of {get_pretty_name(lowerCamelCase		)} to {save_location}'''		)
						torch.save(obj.state_dict()      ,			lowerCamelCase		)

def 		lowerCAmelCase							( lowerCamelCase      ,			lowerCamelCase      ,			lowerCamelCase = 0		):
						'''simple docstring'''


						snake_case__       :					Any										=					f'''{path}/custom_checkpoint_{index}.pkl'''
						logger.info(f'''Loading the state of {get_pretty_name(lowerCamelCase		)} from {load_location}'''		)
						obj.load_state_dict(torch.load(lowerCamelCase      ,			map_location='''cpu'''		)		)


