


"""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


__snake_case         =					False







class _lowerCAmelCase       (   unittest.TestCase			):







   def 				lowerCamelCase	(			self		)			->	List[Any]:

       '''simple docstring'''




       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()







   @property
   def 				lowerCamelCase	(			self		)			->	Dict:

       '''simple docstring'''




       return 12







   @property
   def 				lowerCamelCase	(			self		)			->	Dict:

       '''simple docstring'''




       return 12







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

       '''simple docstring'''




       return 32







   @property
   def 				lowerCamelCase	(			self		)			->	Optional[Any]:

       '''simple docstring'''




       torch.manual_seed(0		)
       snake_case					:   str         =					VQModel(
           block_out_channels=[32, 64]		,			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 				lowerCamelCase	(			self		)			->	Dict:

       '''simple docstring'''




       snake_case					:   str         =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)
       return tokenizer







   @property
   def 				lowerCamelCase	(			self		)			->	str:

       '''simple docstring'''




       torch.manual_seed(0		)
       snake_case					:   Tuple         =					CLIPTextConfig(
           bos_token_id=0		,			eos_token_id=2		,			hidden_size=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=1000		,			)
       return CLIPTextModel(UpperCamelCase__		)







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

       '''simple docstring'''




       torch.manual_seed(0		)

       snake_case					:   Dict         =					12
       snake_case					:   int         =					12

       snake_case					:   Optional[int]         =					{
           "attention_bias": True,
           "cross_attention_dim": 32,
           "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": 32,
           "sample_size": width,
           "activation_fn": "geglu-approximate",
       }

       snake_case					:   str         =					TransformeraDModel(**UpperCamelCase__		)
       return model







   def 				lowerCamelCase	(			self		)			->	Optional[int]:

       '''simple docstring'''




       snake_case					:   Optional[Any]         =					"cpu"

       snake_case					:   Tuple         =					self.dummy_vqvae
       snake_case					:   List[str]         =					self.dummy_text_encoder
       snake_case					:   List[Any]         =					self.dummy_tokenizer
       snake_case					:   Optional[int]         =					self.dummy_transformer
       snake_case					:   List[str]         =					VQDiffusionScheduler(self.num_embed		)
       snake_case					:   List[Any]         =					LearnedClassifierFreeSamplingEmbeddings(learnable=UpperCamelCase__		)

       snake_case					:   List[str]         =					VQDiffusionPipeline(
           vqvae=UpperCamelCase__		,			text_encoder=UpperCamelCase__		,			tokenizer=UpperCamelCase__		,			transformer=UpperCamelCase__		,			scheduler=UpperCamelCase__		,			learned_classifier_free_sampling_embeddings=UpperCamelCase__		,			)
       snake_case					:   int         =					pipe.to(UpperCamelCase__		)
       pipe.set_progress_bar_config(disable=UpperCamelCase__		)

       snake_case					:   Dict         =					"teddy bear playing in the pool"

       snake_case					:   List[Any]         =					torch.Generator(device=UpperCamelCase__		).manual_seed(0		)
       snake_case					:   str         =					pipe([prompt]		,			generator=UpperCamelCase__		,			num_inference_steps=2		,			output_type="np"		)
       snake_case					:   Union[str, Any]         =					output.images

       snake_case					:   Optional[Any]         =					torch.Generator(device=UpperCamelCase__		).manual_seed(0		)
       snake_case					:   Optional[Any]         =					pipe(
           [prompt]		,			generator=UpperCamelCase__		,			output_type="np"		,			return_dict=UpperCamelCase__		,			num_inference_steps=2		)[0]

       snake_case					:   Optional[int]         =					image[0, -3:, -3:, -1]
       snake_case					:   Optional[int]         =					image_from_tuple[0, -3:, -3:, -1]

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

       snake_case					:   Dict         =					np.array([0.6551, 0.6168, 0.5008, 0.5676, 0.5659, 0.4295, 0.6073, 0.5599, 0.4992]		)

       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 				lowerCamelCase	(			self		)			->	Optional[Any]:

       '''simple docstring'''




       snake_case					:   Tuple         =					"cpu"

       snake_case					:   Union[str, Any]         =					self.dummy_vqvae
       snake_case					:   Optional[Any]         =					self.dummy_text_encoder
       snake_case					:   List[Any]         =					self.dummy_tokenizer
       snake_case					:   Union[str, Any]         =					self.dummy_transformer
       snake_case					:   Dict         =					VQDiffusionScheduler(self.num_embed		)
       snake_case					:   List[str]         =					LearnedClassifierFreeSamplingEmbeddings(
           learnable=UpperCamelCase__		,			hidden_size=self.text_embedder_hidden_size		,			length=tokenizer.model_max_length		)

       snake_case					:   str         =					VQDiffusionPipeline(
           vqvae=UpperCamelCase__		,			text_encoder=UpperCamelCase__		,			tokenizer=UpperCamelCase__		,			transformer=UpperCamelCase__		,			scheduler=UpperCamelCase__		,			learned_classifier_free_sampling_embeddings=UpperCamelCase__		,			)
       snake_case					:   Optional[Any]         =					pipe.to(UpperCamelCase__		)
       pipe.set_progress_bar_config(disable=UpperCamelCase__		)

       snake_case					:   Tuple         =					"teddy bear playing in the pool"

       snake_case					:   Optional[Any]         =					torch.Generator(device=UpperCamelCase__		).manual_seed(0		)
       snake_case					:   int         =					pipe([prompt]		,			generator=UpperCamelCase__		,			num_inference_steps=2		,			output_type="np"		)
       snake_case					:   Optional[Any]         =					output.images

       snake_case					:   int         =					torch.Generator(device=UpperCamelCase__		).manual_seed(0		)
       snake_case					:   str         =					pipe(
           [prompt]		,			generator=UpperCamelCase__		,			output_type="np"		,			return_dict=UpperCamelCase__		,			num_inference_steps=2		)[0]

       snake_case					:   Dict         =					image[0, -3:, -3:, -1]
       snake_case					:   List[str]         =					image_from_tuple[0, -3:, -3:, -1]

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

       snake_case					:   Any         =					np.array([0.6693, 0.6075, 0.4959, 0.5701, 0.5583, 0.4333, 0.6171, 0.5684, 0.4988]		)

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







   def 				lowerCamelCase	(			self		)			->	Dict:

       '''simple docstring'''




       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()







   def 				lowerCamelCase	(			self		)			->	List[Any]:

       '''simple docstring'''




       snake_case					:   List[Any]         =					load_numpy(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
           "/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy"		)

       snake_case					:   Tuple         =					VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq"		)
       snake_case					:   Optional[Any]         =					pipeline.to(UpperCamelCase__		)
       pipeline.set_progress_bar_config(disable=UpperCamelCase__		)

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

       snake_case					:   Union[str, Any]         =					output.images[0]

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









"""simple docstring"""



import os
import tempfile
import unittest
from pathlib import Path

from transformers import AutoConfig, is_torch_available
from transformers.testing_utils import require_torch, torch_device


if is_torch_available():
  from transformers import PyTorchBenchmark, PyTorchBenchmarkArguments







@require_torch
class _lowerCAmelCase       (   unittest.TestCase			):







   def 				lowerCamelCase	(			self		,			UpperCamelCase__		)			->	Dict:

       '''simple docstring'''




       for model_result in results.values():
           for batch_size, sequence_length in zip(model_result["bs"]		,			model_result["ss"]		):
               snake_case					:   Optional[int]         =					model_result["result"][batch_size][sequence_length]
               self.assertIsNotNone(UpperCamelCase__		)







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   List[Any]         =					"sshleifer/tiny-gpt2"
       snake_case					:   List[Any]         =					PyTorchBenchmarkArguments(
           models=[MODEL_ID]		,			training=UpperCamelCase__		,			inference=UpperCamelCase__		,			sequence_lengths=[8]		,			batch_sizes=[1]		,			multi_process=UpperCamelCase__		,			)
       snake_case					:   Dict         =					PyTorchBenchmark(UpperCamelCase__		)
       snake_case					:   Union[str, Any]         =					benchmark.run()
       self.check_results_dict_not_empty(results.time_inference_result		)
       self.check_results_dict_not_empty(results.memory_inference_result		)







   def 				lowerCamelCase	(			self		)			->	Union[str, Any]:

       '''simple docstring'''




       snake_case					:   str         =					"sgugger/tiny-distilbert-classification"
       snake_case					:   int         =					PyTorchBenchmarkArguments(
           models=[MODEL_ID]		,			training=UpperCamelCase__		,			inference=UpperCamelCase__		,			sequence_lengths=[8]		,			batch_sizes=[1]		,			multi_process=UpperCamelCase__		,			only_pretrain_model=UpperCamelCase__		,			)
       snake_case					:   Union[str, Any]         =					PyTorchBenchmark(UpperCamelCase__		)
       snake_case					:   Optional[int]         =					benchmark.run()
       self.check_results_dict_not_empty(results.time_inference_result		)
       self.check_results_dict_not_empty(results.memory_inference_result		)







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

       '''simple docstring'''




       snake_case					:   List[Any]         =					"sshleifer/tiny-gpt2"
       snake_case					:   Tuple         =					PyTorchBenchmarkArguments(
           models=[MODEL_ID]		,			training=UpperCamelCase__		,			inference=UpperCamelCase__		,			torchscript=UpperCamelCase__		,			sequence_lengths=[8]		,			batch_sizes=[1]		,			multi_process=UpperCamelCase__		,			)
       snake_case					:   int         =					PyTorchBenchmark(UpperCamelCase__		)
       snake_case					:   int         =					benchmark.run()
       self.check_results_dict_not_empty(results.time_inference_result		)
       self.check_results_dict_not_empty(results.memory_inference_result		)







   @unittest.skipIf(torch_device == "cpu"		,			"Cant do half precision"		)
   def 				lowerCamelCase	(			self		)			->	Optional[Any]:

       '''simple docstring'''




       snake_case					:   List[Any]         =					"sshleifer/tiny-gpt2"
       snake_case					:   int         =					PyTorchBenchmarkArguments(
           models=[MODEL_ID]		,			training=UpperCamelCase__		,			inference=UpperCamelCase__		,			fpaa=UpperCamelCase__		,			sequence_lengths=[8]		,			batch_sizes=[1]		,			multi_process=UpperCamelCase__		,			)
       snake_case					:   List[Any]         =					PyTorchBenchmark(UpperCamelCase__		)
       snake_case					:   str         =					benchmark.run()
       self.check_results_dict_not_empty(results.time_inference_result		)
       self.check_results_dict_not_empty(results.memory_inference_result		)







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

       '''simple docstring'''




       snake_case					:   Optional[int]         =					"sshleifer/tiny-gpt2"
       snake_case					:   int         =					AutoConfig.from_pretrained(UpperCamelCase__		)
       # set architectures equal to `None`
       snake_case					:   str         =					None
       snake_case					:   List[Any]         =					PyTorchBenchmarkArguments(
           models=[MODEL_ID]		,			training=UpperCamelCase__		,			inference=UpperCamelCase__		,			sequence_lengths=[8]		,			batch_sizes=[1]		,			multi_process=UpperCamelCase__		,			)
       snake_case					:   List[Any]         =					PyTorchBenchmark(UpperCamelCase__		,			configs=[config]		)
       snake_case					:   Optional[Any]         =					benchmark.run()
       self.check_results_dict_not_empty(results.time_inference_result		)
       self.check_results_dict_not_empty(results.memory_inference_result		)







   def 				lowerCamelCase	(			self		)			->	Dict:

       '''simple docstring'''




       snake_case					:   Tuple         =					"sshleifer/tiny-gpt2"
       snake_case					:   str         =					PyTorchBenchmarkArguments(
           models=[MODEL_ID]		,			training=UpperCamelCase__		,			inference=UpperCamelCase__		,			sequence_lengths=[8]		,			batch_sizes=[1]		,			multi_process=UpperCamelCase__		,			)
       snake_case					:   List[str]         =					PyTorchBenchmark(UpperCamelCase__		)
       snake_case					:   Dict         =					benchmark.run()
       self.check_results_dict_not_empty(results.time_train_result		)
       self.check_results_dict_not_empty(results.memory_train_result		)







   @unittest.skipIf(torch_device == "cpu"		,			"Can't do half precision"		)
   def 				lowerCamelCase	(			self		)			->	int:

       '''simple docstring'''




       snake_case					:   Optional[Any]         =					"sshleifer/tiny-gpt2"
       snake_case					:   Optional[Any]         =					PyTorchBenchmarkArguments(
           models=[MODEL_ID]		,			training=UpperCamelCase__		,			inference=UpperCamelCase__		,			sequence_lengths=[8]		,			batch_sizes=[1]		,			fpaa=UpperCamelCase__		,			multi_process=UpperCamelCase__		,			)
       snake_case					:   Optional[int]         =					PyTorchBenchmark(UpperCamelCase__		)
       snake_case					:   str         =					benchmark.run()
       self.check_results_dict_not_empty(results.time_train_result		)
       self.check_results_dict_not_empty(results.memory_train_result		)







   def 				lowerCamelCase	(			self		)			->	Optional[int]:

       '''simple docstring'''




       snake_case					:   Tuple         =					"sshleifer/tiny-gpt2"
       snake_case					:   List[Any]         =					AutoConfig.from_pretrained(UpperCamelCase__		)
       snake_case					:   Union[str, Any]         =					PyTorchBenchmarkArguments(
           models=[MODEL_ID]		,			training=UpperCamelCase__		,			inference=UpperCamelCase__		,			sequence_lengths=[8]		,			batch_sizes=[1]		,			multi_process=UpperCamelCase__		,			)
       snake_case					:   Union[str, Any]         =					PyTorchBenchmark(UpperCamelCase__		,			configs=[config]		)
       snake_case					:   int         =					benchmark.run()
       self.check_results_dict_not_empty(results.time_inference_result		)
       self.check_results_dict_not_empty(results.memory_inference_result		)







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   Optional[int]         =					"sshleifer/tinier_bart"
       snake_case					:   List[str]         =					AutoConfig.from_pretrained(UpperCamelCase__		)
       snake_case					:   Optional[int]         =					PyTorchBenchmarkArguments(
           models=[MODEL_ID]		,			training=UpperCamelCase__		,			inference=UpperCamelCase__		,			sequence_lengths=[8]		,			batch_sizes=[1]		,			multi_process=UpperCamelCase__		,			)
       snake_case					:   Union[str, Any]         =					PyTorchBenchmark(UpperCamelCase__		,			configs=[config]		)
       snake_case					:   Optional[Any]         =					benchmark.run()
       self.check_results_dict_not_empty(results.time_inference_result		)
       self.check_results_dict_not_empty(results.memory_inference_result		)







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   Optional[Any]         =					"sshleifer/tiny-gpt2"
       snake_case					:   str         =					AutoConfig.from_pretrained(UpperCamelCase__		)
       snake_case					:   List[str]         =					PyTorchBenchmarkArguments(
           models=[MODEL_ID]		,			training=UpperCamelCase__		,			inference=UpperCamelCase__		,			sequence_lengths=[8]		,			batch_sizes=[1]		,			multi_process=UpperCamelCase__		,			)
       snake_case					:   Optional[int]         =					PyTorchBenchmark(UpperCamelCase__		,			configs=[config]		)
       snake_case					:   int         =					benchmark.run()
       self.check_results_dict_not_empty(results.time_train_result		)
       self.check_results_dict_not_empty(results.memory_train_result		)







   def 				lowerCamelCase	(			self		)			->	Optional[int]:

       '''simple docstring'''




       snake_case					:   str         =					"sshleifer/tinier_bart"
       snake_case					:   Union[str, Any]         =					AutoConfig.from_pretrained(UpperCamelCase__		)
       snake_case					:   Any         =					PyTorchBenchmarkArguments(
           models=[MODEL_ID]		,			training=UpperCamelCase__		,			inference=UpperCamelCase__		,			sequence_lengths=[8]		,			batch_sizes=[1]		,			multi_process=UpperCamelCase__		,			)
       snake_case					:   Union[str, Any]         =					PyTorchBenchmark(UpperCamelCase__		,			configs=[config]		)
       snake_case					:   int         =					benchmark.run()
       self.check_results_dict_not_empty(results.time_train_result		)
       self.check_results_dict_not_empty(results.memory_train_result		)







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   Any         =					"sshleifer/tiny-gpt2"
       with tempfile.TemporaryDirectory() as tmp_dir:
           snake_case					:   Optional[int]         =					PyTorchBenchmarkArguments(
               models=[MODEL_ID]		,			training=UpperCamelCase__		,			inference=UpperCamelCase__		,			save_to_csv=UpperCamelCase__		,			sequence_lengths=[8]		,			batch_sizes=[1]		,			inference_time_csv_file=os.path.join(UpperCamelCase__		,			"inf_time.csv"		)		,			train_memory_csv_file=os.path.join(UpperCamelCase__		,			"train_mem.csv"		)		,			inference_memory_csv_file=os.path.join(UpperCamelCase__		,			"inf_mem.csv"		)		,			train_time_csv_file=os.path.join(UpperCamelCase__		,			"train_time.csv"		)		,			env_info_csv_file=os.path.join(UpperCamelCase__		,			"env.csv"		)		,			multi_process=UpperCamelCase__		,			)
           snake_case					:   List[Any]         =					PyTorchBenchmark(UpperCamelCase__		)
           benchmark.run()
           self.assertTrue(Path(os.path.join(UpperCamelCase__		,			"inf_time.csv"		)		).exists()		)
           self.assertTrue(Path(os.path.join(UpperCamelCase__		,			"train_time.csv"		)		).exists()		)
           self.assertTrue(Path(os.path.join(UpperCamelCase__		,			"inf_mem.csv"		)		).exists()		)
           self.assertTrue(Path(os.path.join(UpperCamelCase__		,			"train_mem.csv"		)		).exists()		)
           self.assertTrue(Path(os.path.join(UpperCamelCase__		,			"env.csv"		)		).exists()		)







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   Union[str, Any]         =					"sshleifer/tiny-gpt2"

       def _check_summary_is_not_empty(UpperCamelCase__		):
           self.assertTrue(hasattr(UpperCamelCase__		,			"sequential"		)		)
           self.assertTrue(hasattr(UpperCamelCase__		,			"cumulative"		)		)
           self.assertTrue(hasattr(UpperCamelCase__		,			"current"		)		)
           self.assertTrue(hasattr(UpperCamelCase__		,			"total"		)		)

       with tempfile.TemporaryDirectory() as tmp_dir:
           snake_case					:   Tuple         =					PyTorchBenchmarkArguments(
               models=[MODEL_ID]		,			training=UpperCamelCase__		,			inference=UpperCamelCase__		,			sequence_lengths=[8]		,			batch_sizes=[1]		,			log_filename=os.path.join(UpperCamelCase__		,			"log.txt"		)		,			log_print=UpperCamelCase__		,			trace_memory_line_by_line=UpperCamelCase__		,			multi_process=UpperCamelCase__		,			)
           snake_case					:   Dict         =					PyTorchBenchmark(UpperCamelCase__		)
           snake_case					:   List[Any]         =					benchmark.run()
           _check_summary_is_not_empty(result.inference_summary		)
           _check_summary_is_not_empty(result.train_summary		)
           self.assertTrue(Path(os.path.join(UpperCamelCase__		,			"log.txt"		)		).exists()		)





