


"""simple docstring"""







import gc
import unittest

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

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


a__				:							Any						=     False




class  UpperCamelCase_	( unittest.TestCase):

	"""simple docstring"""





	def   UpperCAmelCase_      ( self			:				Tuple   )       ->	Optional[int]:
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()




	@property
	def   UpperCAmelCase_      ( self			:				Optional[int]   )       ->	List[Any]:
						return 1_2




	@property
	def   UpperCAmelCase_      ( self			:				List[str]   )       ->	List[str]:
						return 1_2




	@property
	def   UpperCAmelCase_      ( self			:				Dict   )       ->	int:
						return 3_2




	@property
	def   UpperCAmelCase_      ( self			:				Union[str, Any]   )       ->	Optional[int]:
						torch.manual_seed(0   )
						__SCREAMING_SNAKE_CASE						   =				VQModel(
						    block_out_channels=[3_2, 6_4]			,			in_channels=3			,			out_channels=3			,			down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]			,			up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]			,			latent_channels=3			,			num_vq_embeddings=self.num_embed			,			vq_embed_dim=3			,			)
						return model




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




	@property
	def   UpperCAmelCase_      ( self			:				int   )       ->	Any:
						torch.manual_seed(0   )
						__SCREAMING_SNAKE_CASE						   =				CLIPTextConfig(
						    bos_token_id=0			,			eos_token_id=2			,			hidden_size=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 CLIPTextModel(__a   )




	@property
	def   UpperCAmelCase_      ( self			:				Optional[Any]   )       ->	List[str]:
						torch.manual_seed(0   )

						__SCREAMING_SNAKE_CASE						   =				1_2
						__SCREAMING_SNAKE_CASE						   =				1_2

						__SCREAMING_SNAKE_CASE						   =				{
						    "attention_bias": True,
						    "cross_attention_dim": 3_2,
						    "attention_head_dim": height * width,
						    "num_attention_heads": 1,
						    "num_vector_embeds": self.num_embed,
						    "num_embeds_ada_norm": self.num_embeds_ada_norm,
						    "norm_num_groups": 3_2,
						    "sample_size": width,
						    "activation_fn": "geglu-approximate",
						}

						__SCREAMING_SNAKE_CASE						   =				TransformeraDModel(**__a   )
						return model




	def   UpperCAmelCase_      ( self			:				Optional[Any]   )       ->	Tuple:
						__SCREAMING_SNAKE_CASE						   =				"cpu"

						__SCREAMING_SNAKE_CASE						   =				self.dummy_vqvae
						__SCREAMING_SNAKE_CASE						   =				self.dummy_text_encoder
						__SCREAMING_SNAKE_CASE						   =				self.dummy_tokenizer
						__SCREAMING_SNAKE_CASE						   =				self.dummy_transformer
						__SCREAMING_SNAKE_CASE						   =				VQDiffusionScheduler(self.num_embed   )
						__SCREAMING_SNAKE_CASE						   =				LearnedClassifierFreeSamplingEmbeddings(learnable=__a   )

						__SCREAMING_SNAKE_CASE						   =				VQDiffusionPipeline(
						    vqvae=__a			,			text_encoder=__a			,			tokenizer=__a			,			transformer=__a			,			scheduler=__a			,			learned_classifier_free_sampling_embeddings=__a			,			)
						__SCREAMING_SNAKE_CASE						   =				pipe.to(__a   )
						pipe.set_progress_bar_config(disable=__a   )

						__SCREAMING_SNAKE_CASE						   =				"teddy bear playing in the pool"

						__SCREAMING_SNAKE_CASE						   =				torch.Generator(device=__a   ).manual_seed(0   )
						__SCREAMING_SNAKE_CASE						   =				pipe([prompt]			,			generator=__a			,			num_inference_steps=2			,			output_type="np"   )
						__SCREAMING_SNAKE_CASE						   =				output.images

						__SCREAMING_SNAKE_CASE						   =				torch.Generator(device=__a   ).manual_seed(0   )
						__SCREAMING_SNAKE_CASE						   =				pipe(
						    [prompt]			,			generator=__a			,			output_type="np"			,			return_dict=__a			,			num_inference_steps=2   )[0]

						__SCREAMING_SNAKE_CASE						   =				image[0, -3:, -3:, -1]
						__SCREAMING_SNAKE_CASE						   =				image_from_tuple[0, -3:, -3:, -1]

						assert image.shape == (1, 2_4, 2_4, 3)

						__SCREAMING_SNAKE_CASE						   =				np.array([0.6_551, 0.6_168, 0.5_008, 0.5_676, 0.5_659, 0.4_295, 0.6_073, 0.5_599, 0.4_992]   )

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




	def   UpperCAmelCase_      ( self			:				Optional[int]   )       ->	int:
						__SCREAMING_SNAKE_CASE						   =				"cpu"

						__SCREAMING_SNAKE_CASE						   =				self.dummy_vqvae
						__SCREAMING_SNAKE_CASE						   =				self.dummy_text_encoder
						__SCREAMING_SNAKE_CASE						   =				self.dummy_tokenizer
						__SCREAMING_SNAKE_CASE						   =				self.dummy_transformer
						__SCREAMING_SNAKE_CASE						   =				VQDiffusionScheduler(self.num_embed   )
						__SCREAMING_SNAKE_CASE						   =				LearnedClassifierFreeSamplingEmbeddings(
						    learnable=__a			,			hidden_size=self.text_embedder_hidden_size			,			length=tokenizer.model_max_length   )

						__SCREAMING_SNAKE_CASE						   =				VQDiffusionPipeline(
						    vqvae=__a			,			text_encoder=__a			,			tokenizer=__a			,			transformer=__a			,			scheduler=__a			,			learned_classifier_free_sampling_embeddings=__a			,			)
						__SCREAMING_SNAKE_CASE						   =				pipe.to(__a   )
						pipe.set_progress_bar_config(disable=__a   )

						__SCREAMING_SNAKE_CASE						   =				"teddy bear playing in the pool"

						__SCREAMING_SNAKE_CASE						   =				torch.Generator(device=__a   ).manual_seed(0   )
						__SCREAMING_SNAKE_CASE						   =				pipe([prompt]			,			generator=__a			,			num_inference_steps=2			,			output_type="np"   )
						__SCREAMING_SNAKE_CASE						   =				output.images

						__SCREAMING_SNAKE_CASE						   =				torch.Generator(device=__a   ).manual_seed(0   )
						__SCREAMING_SNAKE_CASE						   =				pipe(
						    [prompt]			,			generator=__a			,			output_type="np"			,			return_dict=__a			,			num_inference_steps=2   )[0]

						__SCREAMING_SNAKE_CASE						   =				image[0, -3:, -3:, -1]
						__SCREAMING_SNAKE_CASE						   =				image_from_tuple[0, -3:, -3:, -1]

						assert image.shape == (1, 2_4, 2_4, 3)

						__SCREAMING_SNAKE_CASE						   =				np.array([0.6_693, 0.6_075, 0.4_959, 0.5_701, 0.5_583, 0.4_333, 0.6_171, 0.5_684, 0.4_988]   )

						assert np.abs(image_slice.flatten() - expected_slice   ).max() < 2.0
						assert np.abs(image_from_tuple_slice.flatten() - expected_slice   ).max() < 1E-2




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

	"""simple docstring"""





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




	def   UpperCAmelCase_      ( self			:				int   )       ->	List[str]:
						__SCREAMING_SNAKE_CASE						   =				load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy"   )

						__SCREAMING_SNAKE_CASE						   =				VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq"   )
						__SCREAMING_SNAKE_CASE						   =				pipeline.to(__a   )
						pipeline.set_progress_bar_config(disable=__a   )

						# requires GPU generator for gumbel softmax
						# don't use GPU generator in tests though
						__SCREAMING_SNAKE_CASE						   =				torch.Generator(device=__a   ).manual_seed(0   )
						__SCREAMING_SNAKE_CASE						   =				pipeline(
						    "teddy bear playing in the pool"			,			num_images_per_prompt=1			,			generator=__a			,			output_type="np"			,			)

						__SCREAMING_SNAKE_CASE						   =				output.images[0]

						assert image.shape == (2_5_6, 2_5_6, 3)
						assert np.abs(expected_image - image   ).max() < 2.0





from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_torch_dist_unique_port,
    require_torch_multi_gpu,
    require_torch_neuroncore,
)
from transformers.training_args import ParallelMode
from transformers.utils import logging


_snake_case		=  logging.get_logger(__name__)


if is_torch_available():
						import torch
						from torch import nn
						from torch.utils.data import Dataset

						from transformers import Trainer



						class   UpperCAmelCase_    (       a):







													def __init__(   self,			__a = 101):

																				'''simple docstring'''

																				_lowerCAmelCase     :							str    =    length







													def __len__(   self):

																				'''simple docstring'''

																				return self.length







													def __getitem__(   self,			__a):

																				'''simple docstring'''

																				return i



						class   UpperCAmelCase_    :







													def __call__(   self,			__a):

																				'''simple docstring'''

																				return {"input_ids": torch.tensor(__a), "labels": torch.tensor(__a)}



						class   UpperCAmelCase_    (       nn.Module):







													def __init__(   self):

																				'''simple docstring'''

																				super().__init__()
																				# Add some (unused) params otherwise DDP will complain.
																				_lowerCAmelCase     :							str    =    nn.Linear(120,			80)







													def    snake_case__      (   self,			__a,			__a=None):

																				'''simple docstring'''

																				if labels is not None:
																											return torch.tensor(0.0,			device=input_ids.device), input_ids
																				else:
																											return input_ids



class   UpperCAmelCase_    (       a):







							@require_torch_neuroncore
							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							int    =    f"--nproc_per_node=2\n            --master_port={get_torch_dist_unique_port()}\n            {self.test_file_dir}/test_trainer_distributed.py\n        ".split()
														_lowerCAmelCase     :							Tuple    =    self.get_auto_remove_tmp_dir()
														_lowerCAmelCase     :							Optional[int]    =    f"--output_dir {output_dir}".split()
														_lowerCAmelCase     :							List[Any]    =    ["torchrun"] + distributed_args + args
														execute_subprocess_async(__a,			env=self.get_env())
														# successful return here == success - any errors would have caused an error in the sub-call



class   UpperCAmelCase_    (       a):







							@require_torch_multi_gpu
							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Dict    =    f"--nproc_per_node={torch.cuda.device_count()}\n            --master_port={get_torch_dist_unique_port()}\n            {self.test_file_dir}/test_trainer_distributed.py\n        ".split()
														_lowerCAmelCase     :							Any    =    self.get_auto_remove_tmp_dir()
														_lowerCAmelCase     :							Optional[int]    =    f"--output_dir {output_dir}".split()
														_lowerCAmelCase     :							Any    =    ["torchrun"] + distributed_args + args
														execute_subprocess_async(__a,			env=self.get_env())
														# successful return here == success - any errors would have caused an error in the sub-call


if __name__ == "__main__":
						# The script below is meant to be run under torch.distributed, on a machine with multiple GPUs:
						#
						# PYTHONPATH="src" python -m torch.distributed.run --nproc_per_node 2 --output_dir output_dir ./tests/test_trainer_distributed.py

						_snake_case		=  HfArgumentParser((TrainingArguments,))
						_snake_case		=  parser.parse_args_into_dataclasses()[0]

						logger.warning(
						    f'''Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, '''
						    f'''distributed training: {training_args.parallel_mode != ParallelMode.NOT_DISTRIBUTED}'''
						)

						# Essentially, what we want to verify in the distributed case is that we get all samples back,
						# in the right order. (this is crucial for prediction for instance)
						for dataset_length in [101, 40, 7]:
												_snake_case		=  DummyDataset(dataset_length)






												def 				A				(					_lowerCamelCase  ):





																			'''simple docstring'''




																			_lowerCAmelCase     :							Dict    =    list(range(len(_lowerCamelCase  )  )  )
																			_lowerCAmelCase     :							Union[str, Any]    =    p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
																			if not success and training_args.local_rank == 0:
																										logger.warning(
																										    "Predictions and/or labels do not match expected results:\n  - predictions: "
																										    F"{p.predictions.tolist()}\n  - labels: {p.label_ids.tolist()}\n  - expected: {sequential}"  )
																			return {"success": success}

												_snake_case		=  Trainer(
												    model=DummyModel(),
												    args=training_args,
												    data_collator=DummyDataCollator(),
												    eval_dataset=dataset,
												    compute_metrics=compute_metrics,
												)
												_snake_case		=  trainer.evaluate()
												logger.info(metrics)
												if metrics["eval_success"] is not True:
																		logger.error(metrics)
																		exit(1)

												_snake_case		=  trainer.predict(dataset)
												logger.info(p.metrics)
												if p.metrics["test_success"] is not True:
																		logger.error(p.metrics)
																		exit(1)

												_snake_case		=  2

												_snake_case		=  trainer.evaluate()
												logger.info(metrics)
												if metrics["eval_success"] is not True:
																		logger.error(metrics)
																		exit(1)

												_snake_case		=  trainer.predict(dataset)
												logger.info(p.metrics)
												if p.metrics["test_success"] is not True:
																		logger.error(p.metrics)
																		exit(1)



												_snake_case		=  None
