






import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LDMTextToImagePipeline, UNetaDConditionModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_numpy,
    nightly,
    require_torch_gpu,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class  UpperCamelCase_      (			snake_case_       ,					unittest.TestCase				):



	'''simple docstring'''




	lowerCAmelCase													= LDMTextToImagePipeline
	lowerCAmelCase													= TEXT_TO_IMAGE_PARAMS - {
	    '''negative_prompt''',
	    '''negative_prompt_embeds''',
	    '''cross_attention_kwargs''',
	    '''prompt_embeds''',
	}
	lowerCAmelCase													= PipelineTesterMixin.required_optional_params - {
	    '''num_images_per_prompt''',
	    '''callback''',
	    '''callback_steps''',
	}
	lowerCAmelCase													= TEXT_TO_IMAGE_BATCH_PARAMS
	lowerCAmelCase													= False




	def 							_UpperCamelCase			(       self					)       ->    List[Any]:
								torch.manual_seed(0					)
								snake_case_				=							UNetaDConditionModel(
								    block_out_channels=(32, 64)       ,			layers_per_block=2       ,			sample_size=32       ,			in_channels=4       ,			out_channels=4       ,			down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')       ,			up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')       ,			cross_attention_dim=32       ,			)
								snake_case_				=							DDIMScheduler(
								    beta_start=0.00_085       ,			beta_end=0.012       ,			beta_schedule='scaled_linear'       ,			clip_sample=a       ,			set_alpha_to_one=a       ,			)
								torch.manual_seed(0					)
								snake_case_				=							AutoencoderKL(
								    block_out_channels=(32, 64)       ,			in_channels=3       ,			out_channels=3       ,			down_block_types=('DownEncoderBlock2D', 'DownEncoderBlock2D')       ,			up_block_types=('UpDecoderBlock2D', 'UpDecoderBlock2D')       ,			latent_channels=4       ,			)
								torch.manual_seed(0					)
								snake_case_				=							CLIPTextConfig(
								    bos_token_id=0       ,			eos_token_id=2       ,			hidden_size=32       ,			intermediate_size=37       ,			layer_norm_eps=1E-05       ,			num_attention_heads=4       ,			num_hidden_layers=5       ,			pad_token_id=1       ,			vocab_size=10_00       ,			)
								snake_case_				=							CLIPTextModel(a					)
								snake_case_				=							CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					)

								snake_case_				=							{
								    """unet""": unet,
								    """scheduler""": scheduler,
								    """vqvae""": vae,
								    """bert""": text_encoder,
								    """tokenizer""": tokenizer,
								}
								return components




	def 							_UpperCamelCase			(       self       ,			a       ,			a=0					)       ->    Any:
								if str(a					).startswith('mps'					):
															snake_case_				=							torch.manual_seed(a					)
								else:
															snake_case_				=							torch.Generator(device=a					).manual_seed(a					)
								snake_case_				=							{
								    """prompt""": """A painting of a squirrel eating a burger""",
								    """generator""": generator,
								    """num_inference_steps""": 2,
								    """guidance_scale""": 6.0,
								    """output_type""": """numpy""",
								}
								return inputs






	def 							_UpperCamelCase			(       self					)       ->    Tuple:
								snake_case_				=							"""cpu"""  # ensure determinism for the device-dependent torch.Generator

								snake_case_				=							self.get_dummy_components()
								snake_case_				=							LDMTextToImagePipeline(**a					)
								pipe.to(a					)
								pipe.set_progress_bar_config(disable=a					)

								snake_case_				=							self.get_dummy_inputs(a					)
								snake_case_				=							pipe(**a					).images
								snake_case_				=							image[0, -3:, -3:, -1]

								assert image.shape == (1, 16, 16, 3)
								snake_case_				=							np.array([0.6_101, 0.6_156, 0.5_622, 0.4_895, 0.6_661, 0.3_804, 0.5_748, 0.6_136, 0.5_014]					)

								assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-3






@slow
@require_torch_gpu
class  UpperCamelCase_      (			unittest.TestCase				):



	'''simple docstring'''




	def 							_UpperCamelCase			(       self					)       ->    Optional[Any]:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()




	def 							_UpperCamelCase			(       self       ,			a       ,			a=torch.floataa       ,			a=0					)       ->    Any:
								snake_case_				=							torch.manual_seed(a					)
								snake_case_				=							np.random.RandomState(a					).standard_normal((1, 4, 32, 32)					)
								snake_case_				=							torch.from_numpy(a					).to(device=a       ,			dtype=a					)
								snake_case_				=							{
								    """prompt""": """A painting of a squirrel eating a burger""",
								    """latents""": latents,
								    """generator""": generator,
								    """num_inference_steps""": 3,
								    """guidance_scale""": 6.0,
								    """output_type""": """numpy""",
								}
								return inputs






	def 							_UpperCamelCase			(       self					)       ->    Any:
								snake_case_				=							LDMTextToImagePipeline.from_pretrained('CompVis/ldm-text2im-large-256'					).to(a					)
								pipe.set_progress_bar_config(disable=a					)

								snake_case_				=							self.get_inputs(a					)
								snake_case_				=							pipe(**a					).images
								snake_case_				=							image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 2_56, 2_56, 3)
								snake_case_				=							np.array([0.51_825, 0.52_850, 0.52_543, 0.54_258, 0.52_304, 0.52_569, 0.54_363, 0.55_276, 0.56_878]					)
								snake_case_				=							np.abs(expected_slice - image_slice					).max()
								assert max_diff < 1E-3






@nightly
@require_torch_gpu
class  UpperCamelCase_      (			unittest.TestCase				):



	'''simple docstring'''




	def 							_UpperCamelCase			(       self					)       ->    Dict:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()




	def 							_UpperCamelCase			(       self       ,			a       ,			a=torch.floataa       ,			a=0					)       ->    Any:
								snake_case_				=							torch.manual_seed(a					)
								snake_case_				=							np.random.RandomState(a					).standard_normal((1, 4, 32, 32)					)
								snake_case_				=							torch.from_numpy(a					).to(device=a       ,			dtype=a					)
								snake_case_				=							{
								    """prompt""": """A painting of a squirrel eating a burger""",
								    """latents""": latents,
								    """generator""": generator,
								    """num_inference_steps""": 50,
								    """guidance_scale""": 6.0,
								    """output_type""": """numpy""",
								}
								return inputs






	def 							_UpperCamelCase			(       self					)       ->    Any:
								snake_case_				=							LDMTextToImagePipeline.from_pretrained('CompVis/ldm-text2im-large-256'					).to(a					)
								pipe.set_progress_bar_config(disable=a					)

								snake_case_				=							self.get_inputs(a					)
								snake_case_				=							pipe(**a					).images[0]

								snake_case_				=							load_numpy(
								    'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/ldm_text2img/ldm_large_256_ddim.npy'					)
								snake_case_				=							np.abs(expected_image - image					).max()
								assert max_diff < 1E-3










'''simple docstring'''




# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch

from accelerate import PartialState
from accelerate.utils.operations import broadcast, gather, gather_object, pad_across_processes, reduce



def       __lowerCamelCase  (					_lowercase		)      ->			Tuple:
							return (torch.arange(state.num_processes		) + 1.0 + (state.num_processes * state.process_index)).to(state.device		)



def       __lowerCamelCase  (					_lowercase		)      ->			List[Any]:
							UpperCAmelCase     :		Any							     =		create_tensor(_lowercase		)
							UpperCAmelCase     :		Union[str, Any]							     =		gather(_lowercase		)
							assert gathered_tensor.tolist() == list(range(1	,				state.num_processes**2 + 1		)		)



def       __lowerCamelCase  (					_lowercase		)      ->			Optional[int]:
							UpperCAmelCase     :		Any							     =		[state.process_index]
							UpperCAmelCase     :		Union[str, Any]							     =		gather_object(_lowercase		)
							assert len(_lowercase		) == state.num_processes, F'''{gathered_obj}, {len(_lowercase		)} != {state.num_processes}'''
							assert gathered_obj == list(range(state.num_processes		)		), F'''{gathered_obj} != {list(range(state.num_processes		)		)}'''



def       __lowerCamelCase  (					_lowercase		)      ->			List[Any]:
							UpperCAmelCase     :		Optional[int]							     =		create_tensor(_lowercase		)
							UpperCAmelCase     :		List[str]							     =		broadcast(_lowercase		)
							assert broadcasted_tensor.shape == torch.Size([state.num_processes]		)
							assert broadcasted_tensor.tolist() == list(range(1	,				state.num_processes + 1		)		)



def       __lowerCamelCase  (					_lowercase		)      ->			Tuple:
							# We need to pad the tensor with one more element if we are the main process
							# to ensure that we can pad
							if state.is_main_process:
														UpperCAmelCase     :		Optional[Any]							     =		torch.arange(state.num_processes + 1		).to(state.device		)
							else:
														UpperCAmelCase     :		Tuple							     =		torch.arange(state.num_processes		).to(state.device		)
							UpperCAmelCase     :		Optional[Any]							     =		pad_across_processes(_lowercase		)
							assert padded_tensor.shape == torch.Size([state.num_processes + 1]		)
							if not state.is_main_process:
														assert padded_tensor.tolist() == list(range(0	,				state.num_processes		)		) + [0]



def       __lowerCamelCase  (					_lowercase		)      ->			Dict:
							# For now runs on only two processes
							if state.num_processes != 2:
														return
							UpperCAmelCase     :		Optional[Any]							     =		create_tensor(_lowercase		)
							UpperCAmelCase     :		Optional[Any]							     =		reduce(_lowercase	,				"""sum"""		)
							UpperCAmelCase     :		Optional[Any]							     =		torch.tensor([4.0, 6]		).to(state.device		)
							assert torch.allclose(_lowercase	,				_lowercase		), F'''{reduced_tensor} != {truth_tensor}'''



def       __lowerCamelCase  (					_lowercase		)      ->			Optional[Any]:
							# For now runs on only two processes
							if state.num_processes != 2:
														return
							UpperCAmelCase     :		Tuple							     =		create_tensor(_lowercase		)
							UpperCAmelCase     :		Optional[int]							     =		reduce(_lowercase	,				"""mean"""		)
							UpperCAmelCase     :		str							     =		torch.tensor([2.0, 3]		).to(state.device		)
							assert torch.allclose(_lowercase	,				_lowercase		), F'''{reduced_tensor} != {truth_tensor}'''



def       __lowerCamelCase  (					_lowercase		)      ->			Optional[int]:
							# For xla_spawn (TPUs)
							main()



def       __lowerCamelCase  (					)      ->			int:
							UpperCAmelCase     :		List[Any]							     =		PartialState()
							state.print(F'''State: {state}'''		)
							state.print("""testing gather"""		)
							test_gather(_lowercase		)
							state.print("""testing gather_object"""		)
							test_gather_object(_lowercase		)
							state.print("""testing broadcast"""		)
							test_broadcast(_lowercase		)
							state.print("""testing pad_across_processes"""		)
							test_pad_across_processes(_lowercase		)
							state.print("""testing reduce_sum"""		)
							test_reduce_sum(_lowercase		)
							state.print("""testing reduce_mean"""		)
							test_reduce_mean(_lowercase		)


if __name__ == "__main__":
							main()

