



"""simple docstring"""







import unittest
from functools import partial

import torch

from accelerate import Accelerator, debug_launcher
from accelerate.state import AcceleratorState, GradientState
from accelerate.test_utils import require_cpu, require_huggingface_suite
from accelerate.utils import GradientAccumulationPlugin


def  __SCREAMING_SNAKE_CASE (  __UpperCAmelCase=2   ,				__UpperCAmelCase=True   ,				__UpperCAmelCase=False      ):
		_lowercase				: Dict     =   Accelerator(step_scheduler_with_optimizer=__UpperCAmelCase   ,				split_batches=__UpperCAmelCase      )
		_lowercase				: Union[str, Any]     =   torch.nn.Linear(2   ,				4      )
		_lowercase				: Any     =   torch.optim.AdamW(model.parameters()   ,				lr=1.0      )
		_lowercase				: Optional[int]     =   torch.optim.lr_scheduler.OneCycleLR(__UpperCAmelCase   ,				max_lr=0.0_1   ,				steps_per_epoch=2   ,				epochs=1      )
		_lowercase     ,  _lowercase     ,  _lowercase				: Any     =   accelerator.prepare(__UpperCAmelCase   ,				__UpperCAmelCase   ,				__UpperCAmelCase      )

		# Optimizer has stepped
		scheduler.step()
		if step_scheduler_with_optimizer or (num_processes == 1):
				assert (
				    scheduler.scheduler.last_epoch == num_processes
				), F"""Last Epoch ({scheduler.scheduler.last_epoch}) != Num Processes ({num_processes})"""
		else:
				assert (
				    scheduler.scheduler.last_epoch != num_processes
				), F"""Last Epoch ({scheduler.scheduler.last_epoch}) == Num Processes ({num_processes})"""


def  __SCREAMING_SNAKE_CASE (  __UpperCAmelCase=2   ,				__UpperCAmelCase=True   ,				__UpperCAmelCase=False      ):
		_lowercase				: str     =   Accelerator(step_scheduler_with_optimizer=__UpperCAmelCase   ,				split_batches=__UpperCAmelCase      )
		_lowercase				: int     =   torch.nn.Linear(2   ,				4      )
		_lowercase				: int     =   torch.optim.AdamW(model.parameters()   ,				lr=1.0      )
		_lowercase				: int     =   torch.optim.lr_scheduler.LambdaLR(__UpperCAmelCase   ,				lr_lambda=lambda __UpperCAmelCase      : 1 - n / 10      )
		_lowercase     ,  _lowercase     ,  _lowercase				: List[Any]     =   accelerator.prepare(__UpperCAmelCase   ,				__UpperCAmelCase   ,				__UpperCAmelCase      )

		# Optimizer has stepped
		_lowercase				: List[Any]     =   False
		scheduler.step()
		_lowercase				: List[str]     =   1 - (num_processes if (step_scheduler_with_optimizer and not split_batches) else 1) / 10
		assert (
		    scheduler.get_last_lr()[0] == expected_lr
		), F"""Wrong lr found at first step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}"""

		# Optimizer has not stepped
		_lowercase				: Optional[Any]     =   True
		scheduler.step()
		if not step_scheduler_with_optimizer:
				_lowercase				: Tuple     =   1 - 2 / 10
		assert (
		    scheduler.get_last_lr()[0] == expected_lr
		), F"""Wrong lr found at second step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}"""


def  __SCREAMING_SNAKE_CASE (  __UpperCAmelCase = 2      ):
		from transformers import get_linear_schedule_with_warmup

		_lowercase				: Tuple     =   [1, 2, 4]
		for num_steps in steps:
				_lowercase				: Any     =   GradientAccumulationPlugin(num_steps=__UpperCAmelCase   ,				adjust_scheduler=num_steps > 1      )
				_lowercase				: int     =   Accelerator(gradient_accumulation_plugin=__UpperCAmelCase      )
				_lowercase				: int     =   torch.nn.Linear(2   ,				4      )
				_lowercase				: int     =   torch.optim.AdamW(model.parameters()   ,				lr=1_0.0      )
				_lowercase				: Dict     =   get_linear_schedule_with_warmup(optimizer=__UpperCAmelCase   ,				num_warmup_steps=0   ,				num_training_steps=20      )

				_lowercase     ,  _lowercase     ,  _lowercase				: List[str]     =   accelerator.prepare(__UpperCAmelCase   ,				__UpperCAmelCase   ,				__UpperCAmelCase      )

				for i in range(10 * num_steps      ):
						with accelerator.accumulate(__UpperCAmelCase      ):
								optimizer.step()
								scheduler.step()

						if i == (10 * num_steps - 2):
								assert (
								    scheduler.get_last_lr()[0] != 0
								), F"""Wrong lr found at second-to-last step, expected non-zero, got {scheduler.get_last_lr()[0]}. num_steps: {num_steps}"""
				assert (
				    scheduler.get_last_lr()[0] == 0
				), F"""Wrong lr found at last step, expected 0, got {scheduler.get_last_lr()[0]}"""
				GradientState._reset_state()





@require_cpu
class   UpperCamelCase	(     unittest.TestCase       ):




	"""simple docstring"""






	def    lowerCamelCase__					(     self   ):
			debug_launcher(partial(UpperCAmelCase_				,num_processes=1   )				,num_processes=1   )
			debug_launcher(partial(UpperCAmelCase_				,num_processes=1				,split_batches=UpperCAmelCase_   )				,num_processes=1   )






	def    lowerCamelCase__					(     self   ):
			debug_launcher(partial(UpperCAmelCase_				,num_processes=1   )				,num_processes=1   )
			debug_launcher(partial(UpperCAmelCase_				,num_processes=1				,split_batches=UpperCAmelCase_   )				,num_processes=1   )






	def    lowerCamelCase__					(     self   ):
			debug_launcher(partial(UpperCAmelCase_				,num_processes=1				,step_scheduler_with_optimizer=UpperCAmelCase_   )				,num_processes=1   )






	def    lowerCamelCase__					(     self   ):
			debug_launcher(partial(UpperCAmelCase_				,num_processes=1				,step_scheduler_with_optimizer=UpperCAmelCase_   )				,num_processes=1   )






	def    lowerCamelCase__					(     self   ):
			AcceleratorState._reset_state(UpperCAmelCase_   )
			debug_launcher(UpperCAmelCase_   )
			debug_launcher(partial(UpperCAmelCase_				,num_processes=1				,split_batches=UpperCAmelCase_   )				,num_processes=1   )






	def    lowerCamelCase__					(     self   ):
			AcceleratorState._reset_state(UpperCAmelCase_   )
			debug_launcher(UpperCAmelCase_   )
			debug_launcher(partial(UpperCAmelCase_				,num_processes=1				,split_batches=UpperCAmelCase_   )				,num_processes=1   )






	def    lowerCamelCase__					(     self   ):
			AcceleratorState._reset_state(UpperCAmelCase_   )
			debug_launcher(partial(UpperCAmelCase_				,step_scheduler_with_optimizer=UpperCAmelCase_   )   )






	def    lowerCamelCase__					(     self   ):
			AcceleratorState._reset_state(UpperCAmelCase_   )
			debug_launcher(partial(UpperCAmelCase_				,step_scheduler_with_optimizer=UpperCAmelCase_   )   )






	@require_huggingface_suite
	def    lowerCamelCase__					(     self   ):
			AcceleratorState._reset_state(UpperCAmelCase_   )
			debug_launcher(partial(UpperCAmelCase_				,num_processes=1   )   )
			debug_launcher(UpperCAmelCase_   )




"""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 							snake_case			(					__UpperCAmelCase ,       unittest.TestCase							):




					"""simple docstring"""

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







					@property
					def 						__lowerCAmelCase		(							self  :      Tuple     ):
									return 32







					@property
					def 						__lowerCAmelCase		(							self  :      Union[str, Any]     ):
									return 32







					@property
					def 						__lowerCAmelCase		(							self  :      str     ):
									return self.time_input_dim







					@property
					def 						__lowerCAmelCase		(							self  :      Union[str, Any]     ):
									return self.time_input_dim * 4







					@property
					def 						__lowerCAmelCase		(							self  :      Optional[Any]     ):
									return 100







					@property
					def 						__lowerCAmelCase		(							self  :      int     ):
									UpperCAmelCase__			 =				CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )
									return tokenizer







					@property
					def 						__lowerCAmelCase		(							self  :      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=1_000       ,)
									return CLIPTextModelWithProjection(lowerCamelCase__     )







					@property
					def 						__lowerCAmelCase		(							self  :      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 						__lowerCAmelCase		(							self  :      Optional[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 						__lowerCAmelCase		(							self  :      List[str]     ):
									UpperCAmelCase__			 =				CLIPImageProcessor(
									    crop_size=224       ,do_center_crop=lowerCamelCase__       ,do_normalize=lowerCamelCase__       ,do_resize=lowerCamelCase__       ,image_mean=[0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3]       ,image_std=[0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1]       ,resample=3       ,size=224       ,)

									return image_processor







					def 						__lowerCAmelCase		(							self  :      List[Any]     ):
									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=1_000       ,clip_sample=lowerCamelCase__       ,clip_sample_range=1_0.0       ,)

									UpperCAmelCase__			 =				{
									    '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]       ,lowerCamelCase__  :      Any       ,lowerCamelCase__  :      str=0     ):
									if str(lowerCamelCase__     ).startswith('mps'     ):
													UpperCAmelCase__			 =				torch.manual_seed(lowerCamelCase__     )
									else:
													UpperCAmelCase__			 =				torch.Generator(device=lowerCamelCase__     ).manual_seed(lowerCamelCase__     )

									UpperCAmelCase__			 =				floats_tensor((1, 3, 64, 64)       ,rng=random.Random(lowerCamelCase__     )     ).to(lowerCamelCase__     )
									UpperCAmelCase__			 =				image.cpu().permute(0       ,2       ,3       ,1     )[0]
									UpperCAmelCase__			 =				Image.fromarray(np.uinta(lowerCamelCase__     )     ).convert('RGB'     ).resize((256, 256)     )

									UpperCAmelCase__			 =				{
									    '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  :      List[Any]     ):
									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.1_0_7_1_2_8_4,
									        1.3_3_3_0_2_7_1,
									        0.6_1_2_6_0_2_2_3,
									        -0.6_6_9_1_0_6_5,
									        -0.3_8_4_6_8_5_2,
									        -1.0_3_0_3_6_6_1,
									        0.2_2_7_1_6_1_1_1,
									        0.0_3_3_4_8_9_0_1,
									        0.3_0_0_4_0_6_7_5,
									        -0.2_4_8_0_5_0_2_9,
									    ]     )

									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  :      Optional[int]     ):
									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 						__lowerCAmelCase		(							self  :      int     ):
									UpperCAmelCase__			 =				torch_device == 'cpu'
									UpperCAmelCase__			 =				False

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



