



from copy import deepcopy

import torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader

from accelerate.accelerator import Accelerator
from accelerate.state import GradientState
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import DistributedType, is_torch_version, set_seed






def      lowerCAmelCase__		(   lowerCamelCase_		:		int	,lowerCamelCase_		:		Tuple	,lowerCamelCase_		:		Any	,lowerCamelCase_		:		str):





				'''simple docstring'''

				for param, grad_param in zip(model_a.parameters()	,model_b.parameters()):
								if not param.requires_grad:
												continue
								if not did_step:
												# Grads should not be in sync
												assert (
												    torch.allclose(param.grad	,grad_param.grad) is False
												), f"""Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})"""
								else:
												# Grads should be in sync
												assert (
												    torch.allclose(param.grad	,grad_param.grad) is True
												), f"""Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})"""






def      lowerCAmelCase__		(   lowerCamelCase_		:		Optional[Any]	,lowerCamelCase_		:		int	,lowerCamelCase_		:		Union[str, Any]	,lowerCamelCase_		:		Dict	,lowerCamelCase_		:		int=True):





				'''simple docstring'''

				model.train()
				lowerCAmelCase__					:    Union[str, Any]												=					model(lowerCamelCase_)
				lowerCAmelCase__					:    Optional[Any]												=					F.mse_loss(lowerCamelCase_	,target.to(output.device))
				if not do_backward:
								loss /= accelerator.gradient_accumulation_steps
								loss.backward()
				else:
								accelerator.backward(lowerCamelCase_)






def      lowerCAmelCase__		(   lowerCamelCase_		:		List[Any]	,lowerCamelCase_		:		Tuple=False):





				'''simple docstring'''

				set_seed(42)
				lowerCAmelCase__					:    Tuple												=					RegressionModel()
				lowerCAmelCase__					:    Optional[int]												=					deepcopy(lowerCamelCase_)
				lowerCAmelCase__					:    int												=					RegressionDataset(length=80)
				lowerCAmelCase__					:    int												=					DataLoader(lowerCamelCase_	,batch_size=16)
				model.to(accelerator.device)
				if sched:
								lowerCAmelCase__					:    int												=					AdamW(params=model.parameters()	,lr=1E-3)
								lowerCAmelCase__					:    Dict												=					AdamW(params=ddp_model.parameters()	,lr=1E-3)
								lowerCAmelCase__					:    List[str]												=					LambdaLR(lowerCamelCase_	,lr_lambda=lambda lowerCamelCase_: epoch**0.65)
								lowerCAmelCase__					:    Dict												=					LambdaLR(lowerCamelCase_	,lr_lambda=lambda lowerCamelCase_: epoch**0.65)
				# Make a copy of `model`
				if sched:
								lowerCAmelCase__					:    List[str]												=					accelerator.prepare(lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_)
				else:
								lowerCAmelCase__					:    int												=					accelerator.prepare(lowerCamelCase_	,lowerCamelCase_)
				if sched:
								return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
				return model, ddp_model, dataloader






def      lowerCAmelCase__		(   lowerCamelCase_		:		Tuple):





				'''simple docstring'''

				lowerCAmelCase__					:    str												=					get_training_setup(lowerCamelCase_)
				# Use a single batch
				lowerCAmelCase__					:    List[Any]												=					next(iter(lowerCamelCase_)).values()
				for iteration in range(3):
								# Gather the distributed inputs and targs for the base model
								lowerCAmelCase__					:    List[Any]												=					accelerator.gather((ddp_input, ddp_target))
								lowerCAmelCase__					:    str												=					input.to(accelerator.device), target.to(accelerator.device)
								# Perform our initial ground truth step in non "DDP"
								step_model(lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_)
								# Do "gradient accumulation" (noop)
								if iteration % 2 == 0:
												# Accumulate grads locally
												with accelerator.no_sync(lowerCamelCase_):
																step_model(lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_)
								else:
												# Sync grads
												step_model(lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_)

								# Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
								check_model_parameters(lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_)
								for param, ddp_param in zip(model.parameters()	,ddp_model.parameters()):
												if not param.requires_grad:
																continue
												assert torch.allclose(
												    param.grad	,ddp_param.grad), f"""Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"""

								# Shuffle ddp_input on each iteration
								torch.manual_seed(1337 + iteration)
								lowerCAmelCase__					:    List[str]												=					ddp_input[torch.randperm(len(lowerCamelCase_))]






def      lowerCAmelCase__		(   lowerCamelCase_		:		Tuple):





				'''simple docstring'''

				lowerCAmelCase__					:    List[str]												=					get_training_setup(lowerCamelCase_)
				# Use a single batch
				lowerCAmelCase__					:    int												=					next(iter(lowerCamelCase_)).values()
				for iteration in range(3):
								# Gather the distributed inputs and targs for the base model
								lowerCAmelCase__					:    List[str]												=					accelerator.gather((ddp_input, ddp_target))
								lowerCAmelCase__					:    List[str]												=					input.to(accelerator.device), target.to(accelerator.device)
								# Perform our initial ground truth step in non "DDP"
								step_model(lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_)
								# Do "gradient accumulation" (noop)
								if iteration % 2 == 0:
												# Accumulate grads locally
												with accelerator.no_sync(lowerCamelCase_):
																step_model(lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_)
								else:
												# Sync grads
												step_model(lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_)

								# DDP model and model should only be in sync when not (iteration % 2 == 0)
								for param, ddp_param in zip(model.parameters()	,ddp_model.parameters()):
												if not param.requires_grad:
																continue
												if iteration % 2 == 0:
																# Grads should not be in sync
																assert (
																    torch.allclose(param.grad	,ddp_param.grad) is False
																), f"""Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"""
												else:
																# Grads should be in sync
																assert (
																    torch.allclose(param.grad	,ddp_param.grad) is True
																), f"""Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"""

        # Shuffle ddp_input on each iteration
								torch.manual_seed(1337 + iteration)
								lowerCAmelCase__					:    str												=					ddp_input[torch.randperm(len(lowerCamelCase_))]






def      lowerCAmelCase__		(   lowerCamelCase_		:		Dict=False	,lowerCamelCase_		:		Any=False):





				'''simple docstring'''

				lowerCAmelCase__					:    List[str]												=					Accelerator(
				    split_batches=lowerCamelCase_	,dispatch_batches=lowerCamelCase_	,gradient_accumulation_steps=2)
				# Test that context manager behaves properly
				lowerCAmelCase__					:    int												=					get_training_setup(lowerCamelCase_)
				for iteration, batch in enumerate(lowerCamelCase_):
								lowerCAmelCase__					:    Optional[Any]												=					batch.values()
								# Gather the distributed inputs and targs for the base model
								lowerCAmelCase__					:    str												=					accelerator.gather((ddp_input, ddp_target))
								lowerCAmelCase__					:    Tuple												=					input.to(accelerator.device), target.to(accelerator.device)
								# Perform our initial ground truth step in non "DDP"
								step_model(lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_)
								# Do "gradient accumulation" (noop)
								with accelerator.accumulate(lowerCamelCase_):
												step_model(lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_)

								# DDP model and model should only be in sync when not (iteration % 2 == 0)
								for param, ddp_param in zip(model.parameters()	,ddp_model.parameters()):
												if not param.requires_grad:
																continue
												if ((iteration + 1) % 2 == 0) or (iteration == len(lowerCamelCase_) - 1):
																# Grads should be in sync
																assert (
																    torch.allclose(param.grad	,ddp_param.grad) is True
																), f"""Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"""
												else:
																# Grads should not be in sync
																assert (
																    torch.allclose(param.grad	,ddp_param.grad) is False
																), f"""Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"""

        # Shuffle ddp_input on each iteration
								torch.manual_seed(1337 + iteration)
								lowerCAmelCase__					:    Dict												=					ddp_input[torch.randperm(len(lowerCamelCase_))]
				GradientState._reset_state()






def      lowerCAmelCase__		(   lowerCamelCase_		:		int=False	,lowerCamelCase_		:		Optional[int]=False):





				'''simple docstring'''

				lowerCAmelCase__					:    Optional[int]												=					Accelerator(
				    split_batches=lowerCamelCase_	,dispatch_batches=lowerCamelCase_	,gradient_accumulation_steps=2)
				# Test that context manager behaves properly
				lowerCAmelCase__					:    Optional[int]												=					get_training_setup(lowerCamelCase_	,lowerCamelCase_)
				for iteration, batch in enumerate(lowerCamelCase_):
								lowerCAmelCase__					:    Optional[Any]												=					batch.values()
								# Gather the distributed inputs and targs for the base model
								lowerCAmelCase__					:    Tuple												=					accelerator.gather((ddp_input, ddp_target))
								lowerCAmelCase__					:    str												=					input.to(accelerator.device), target.to(accelerator.device)
								# Perform our initial ground truth step in non "DDP"
								model.train()
								ddp_model.train()
								step_model(lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_)
								opt.step()

								if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(lowerCamelCase_)):
												if split_batches:
																sched.step()
												else:
																for _ in range(accelerator.num_processes):
																				sched.step()
								opt.zero_grad()
								# Perform gradient accumulation under wrapper
								with accelerator.accumulate(lowerCamelCase_):
												step_model(lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_)
												ddp_opt.step()
												ddp_sched.step()
												ddp_opt.zero_grad()

								# Learning rates should be the same
								assert (
								    opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
								), f"""Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]['lr']}\nDDP opt: {ddp_opt.param_groups[0]['lr']}\n"""
								lowerCAmelCase__					:    List[Any]												=					(((iteration + 1) % 2) == 0) or ((iteration + 1) == len(lowerCamelCase_))
								if accelerator.num_processes > 1:
												check_model_parameters(lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_	,lowerCamelCase_)
								# Shuffle ddp_input on each iteration
								torch.manual_seed(1337 + iteration)
				GradientState._reset_state()






def      lowerCAmelCase__		(   ):





				'''simple docstring'''

				lowerCAmelCase__					:    int												=					Accelerator()

				lowerCAmelCase__					:    Optional[int]												=					RegressionDataset(length=80)
				lowerCAmelCase__					:    List[Any]												=					DataLoader(lowerCamelCase_	,batch_size=16)
				lowerCAmelCase__					:    Union[str, Any]												=					RegressionDataset(length=96)
				lowerCAmelCase__					:    Optional[Any]												=					DataLoader(lowerCamelCase_	,batch_size=16)
				lowerCAmelCase__					:    Optional[Any]												=					accelerator.prepare(lowerCamelCase_	,lowerCamelCase_)
				assert accelerator.gradient_state.active_dataloader is None
				for iteration, _ in enumerate(lowerCamelCase_):
								assert id(accelerator.gradient_state.active_dataloader) == id(lowerCamelCase_)
								if iteration < len(lowerCamelCase_) - 1:
												assert not accelerator.gradient_state.end_of_dataloader
												if iteration == 1:
																for batch_num, _ in enumerate(lowerCamelCase_):
																				assert id(accelerator.gradient_state.active_dataloader) == id(lowerCamelCase_)
																				if batch_num < len(lowerCamelCase_) - 1:
																								assert not accelerator.gradient_state.end_of_dataloader
																				else:
																								assert accelerator.gradient_state.end_of_dataloader
								else:
												assert accelerator.gradient_state.end_of_dataloader
				assert accelerator.gradient_state.active_dataloader is None






def      lowerCAmelCase__		(   ):





				'''simple docstring'''

				lowerCAmelCase__					:    List[Any]												=					Accelerator()
				lowerCAmelCase__					:    int												=					accelerator.state
				if state.local_process_index == 0:
								print('''**Test `accumulate` gradient accumulation with dataloader break**''')
				test_dataloader_break()
				if state.distributed_type == DistributedType.NO:
								if state.local_process_index == 0:
												print('''**Test NOOP `no_sync` context manager**''')
								test_noop_sync(lowerCamelCase_)
				if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
								if state.local_process_index == 0:
												print('''**Test Distributed `no_sync` context manager**''')
								test_distributed_sync(lowerCamelCase_)
				if state.distributed_type == DistributedType.MULTI_GPU:
								for split_batch in [True, False]:
												for dispatch_batches in [True, False]:
																if state.local_process_index == 0:
																				print(
																				    '''**Test `accumulate` gradient accumulation, '''	,f"""`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**"""	,)
																test_gradient_accumulation(lowerCamelCase_	,lowerCamelCase_)

    # Currently will break on torch 2.0 +, need to investigate why
				if is_torch_version('''<'''	,'''2.0''') or state.distributed_type == DistributedType.NO:
								if state.local_process_index == 0:
												print(
												    '''**Test `accumulate` gradient accumulation with optimizer and scheduler, '''	,'''`split_batches=False`, `dispatch_batches=False`**'''	,)
								test_gradient_accumulation_with_opt_and_scheduler()
								if state.distributed_type == DistributedType.MULTI_GPU:
												for split_batch in [True, False]:
																for dispatch_batches in [True, False]:
																				if not split_batch and not dispatch_batches:
																								continue
																				if state.local_process_index == 0:
																								print(
																								    '''**Test `accumulate` gradient accumulation with optimizer and scheduler, '''	,f"""`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**"""	,)
																				test_gradient_accumulation_with_opt_and_scheduler(lowerCamelCase_	,lowerCamelCase_)






def      lowerCAmelCase__		(   lowerCamelCase_		:		Any):





				'''simple docstring'''

				main()


if __name__ == "__main__":
							main()






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 	lowerCamelCase__      (							lowerCamelCase__				,    unittest.TestCase):




		'''simple docstring'''



		snake_case_     =KandinskyVaaPriorPipeline
		snake_case_     =["""prompt"""]
		snake_case_     =["""prompt""", """negative_prompt"""]
		snake_case_     =[
		    """num_images_per_prompt""",
		    """generator""",
		    """num_inference_steps""",
		    """latents""",
		    """negative_prompt""",
		    """guidance_scale""",
		    """output_type""",
		    """return_dict""",
		]
		snake_case_     =False



		@property
		def 		lowerCAmelCase__				(self			)					->      Any:

						"""simple docstring"""






						return 32



		@property
		def 		lowerCAmelCase__				(self			)					->      List[str]:

						"""simple docstring"""






						return 32



		@property
		def 		lowerCAmelCase__				(self			)					->      Tuple:

						"""simple docstring"""






						return self.time_input_dim



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

						"""simple docstring"""






						return self.time_input_dim * 4



		@property
		def 		lowerCAmelCase__				(self			)					->      Dict:

						"""simple docstring"""






						return 1_00



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

						"""simple docstring"""






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



		@property
		def 		lowerCAmelCase__				(self			)					->      str:

						"""simple docstring"""






						torch.manual_seed(0			)
						lowerCAmelCase__					:    str												=					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=10_00		,)
						return CLIPTextModelWithProjection(__lowerCamelCase			)



		@property
		def 		lowerCAmelCase__				(self			)					->      int:

						"""simple docstring"""






						torch.manual_seed(0			)

						lowerCAmelCase__					:    List[str]												=					{
						    '''num_attention_heads''': 2,
						    '''attention_head_dim''': 12,
						    '''embedding_dim''': self.text_embedder_hidden_size,
						    '''num_layers''': 1,
						}

						lowerCAmelCase__					:    List[str]												=					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
						lowerCAmelCase__					:    int												=					nn.Parameter(torch.ones(model.clip_std.shape			)			)
						return model



		@property
		def 		lowerCAmelCase__				(self			)					->      List[str]:

						"""simple docstring"""






						torch.manual_seed(0			)
						lowerCAmelCase__					:    Any												=					CLIPVisionConfig(
						    hidden_size=self.text_embedder_hidden_size		,image_size=2_24		,projection_dim=self.text_embedder_hidden_size		,intermediate_size=37		,num_attention_heads=4		,num_channels=3		,num_hidden_layers=5		,patch_size=14		,)

						lowerCAmelCase__					:    int												=					CLIPVisionModelWithProjection(__lowerCamelCase			)
						return model



		@property
		def 		lowerCAmelCase__				(self			)					->      Any:

						"""simple docstring"""






						lowerCAmelCase__					:    Optional[Any]												=					CLIPImageProcessor(
						    crop_size=2_24		,do_center_crop=__lowerCamelCase		,do_normalize=__lowerCamelCase		,do_resize=__lowerCamelCase		,image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]		,image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]		,resample=3		,size=2_24		,)

						return image_processor



		def 		lowerCAmelCase__				(self			)					->      Any:

						"""simple docstring"""






						lowerCAmelCase__					:    List[Any]												=					self.dummy_prior
						lowerCAmelCase__					:    List[Any]												=					self.dummy_image_encoder
						lowerCAmelCase__					:    Optional[Any]												=					self.dummy_text_encoder
						lowerCAmelCase__					:    Optional[int]												=					self.dummy_tokenizer
						lowerCAmelCase__					:    str												=					self.dummy_image_processor

						lowerCAmelCase__					:    Union[str, Any]												=					UnCLIPScheduler(
						    variance_type='''fixed_small_log'''		,prediction_type='''sample'''		,num_train_timesteps=10_00		,clip_sample=__lowerCamelCase		,clip_sample_range=10.0		,)

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

						return components



		def 		lowerCAmelCase__				(self		,__lowerCamelCase		,__lowerCamelCase=0			)					->      int:

						"""simple docstring"""






						if str(__lowerCamelCase			).startswith('''mps'''			):
										lowerCAmelCase__					:    Dict												=					torch.manual_seed(__lowerCamelCase			)
						else:
										lowerCAmelCase__					:    Any												=					torch.Generator(device=__lowerCamelCase			).manual_seed(__lowerCamelCase			)
						lowerCAmelCase__					:    Optional[Any]												=					{
						    '''prompt''': '''horse''',
						    '''generator''': generator,
						    '''guidance_scale''': 4.0,
						    '''num_inference_steps''': 2,
						    '''output_type''': '''np''',
						}
						return inputs



		def 		lowerCAmelCase__				(self			)					->      List[str]:

						"""simple docstring"""






						lowerCAmelCase__					:    str												=					'''cpu'''

						lowerCAmelCase__					:    Optional[int]												=					self.get_dummy_components()

						lowerCAmelCase__					:    int												=					self.pipeline_class(**__lowerCamelCase			)
						lowerCAmelCase__					:    int												=					pipe.to(__lowerCamelCase			)

						pipe.set_progress_bar_config(disable=__lowerCamelCase			)

						lowerCAmelCase__					:    List[Any]												=					pipe(**self.get_dummy_inputs(__lowerCamelCase			)			)
						lowerCAmelCase__					:    Union[str, Any]												=					output.image_embeds

						lowerCAmelCase__					:    Tuple												=					pipe(
						    **self.get_dummy_inputs(__lowerCamelCase			)		,return_dict=__lowerCamelCase		,)[0]

						lowerCAmelCase__					:    Union[str, Any]												=					image[0, -10:]
						lowerCAmelCase__					:    str												=					image_from_tuple[0, -10:]

						assert image.shape == (1, 32)

						lowerCAmelCase__					:    int												=					np.array(
						    [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156]			)

						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]:

						"""simple docstring"""






						lowerCAmelCase__					:    List[Any]												=					torch_device == '''cpu'''
						lowerCAmelCase__					:    List[Any]												=					True
						lowerCAmelCase__					:    Optional[int]												=					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			)					->      Any:

						"""simple docstring"""






						lowerCAmelCase__					:    int												=					torch_device == '''cpu'''
						lowerCAmelCase__					:    int												=					False

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

