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


UpperCAmelCase        =				False



class 			lowercase					( unittest.TestCase		):






			def 					UpperCAmelCase							(self  :	int			)       ->	Dict:





				"""simple docstring"""


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






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





				"""simple docstring"""


				return 12






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





				"""simple docstring"""


				return 12






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





				"""simple docstring"""


				return 32






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





				"""simple docstring"""


				torch.manual_seed(0			)
				lowerCAmelCase												=  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 					UpperCAmelCase							(self  :	int			)       ->	int:





				"""simple docstring"""


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






			@property
			def 					UpperCAmelCase							(self  :	Any			)       ->	List[Any]:





				"""simple docstring"""


				torch.manual_seed(0			)
				lowerCAmelCase												=  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=1_000      ,)
				return CLIPTextModel(SCREAMING_SNAKE_CASE_			)






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





				"""simple docstring"""


				torch.manual_seed(0			)

				lowerCAmelCase												=  12
				lowerCAmelCase												=  12

				lowerCAmelCase												=  {
				    '''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''',
				}

				lowerCAmelCase												=  TransformeraDModel(**SCREAMING_SNAKE_CASE_			)
				return model






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





				"""simple docstring"""


				lowerCAmelCase												=  '''cpu'''

				lowerCAmelCase												=  self.dummy_vqvae
				lowerCAmelCase												=  self.dummy_text_encoder
				lowerCAmelCase												=  self.dummy_tokenizer
				lowerCAmelCase												=  self.dummy_transformer
				lowerCAmelCase												=  VQDiffusionScheduler(self.num_embed			)
				lowerCAmelCase												=  LearnedClassifierFreeSamplingEmbeddings(learnable=SCREAMING_SNAKE_CASE_			)

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

				lowerCAmelCase												=  '''teddy bear playing in the pool'''

				lowerCAmelCase												=  torch.Generator(device=SCREAMING_SNAKE_CASE_			).manual_seed(0			)
				lowerCAmelCase												=  pipe([prompt]      ,generator=SCREAMING_SNAKE_CASE_      ,num_inference_steps=2      ,output_type='''np'''			)
				lowerCAmelCase												=  output.images

				lowerCAmelCase												=  torch.Generator(device=SCREAMING_SNAKE_CASE_			).manual_seed(0			)
				lowerCAmelCase												=  pipe(
				    [prompt]      ,generator=SCREAMING_SNAKE_CASE_      ,output_type='''np'''      ,return_dict=SCREAMING_SNAKE_CASE_      ,num_inference_steps=2			)[0]

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

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

				lowerCAmelCase												=  np.array([0.65_51, 0.61_68, 0.50_08, 0.56_76, 0.56_59, 0.42_95, 0.60_73, 0.55_99, 0.49_92]			)

				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  :	List[str]			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase												=  '''cpu'''

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

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

				lowerCAmelCase												=  '''teddy bear playing in the pool'''

				lowerCAmelCase												=  torch.Generator(device=SCREAMING_SNAKE_CASE_			).manual_seed(0			)
				lowerCAmelCase												=  pipe([prompt]      ,generator=SCREAMING_SNAKE_CASE_      ,num_inference_steps=2      ,output_type='''np'''			)
				lowerCAmelCase												=  output.images

				lowerCAmelCase												=  torch.Generator(device=SCREAMING_SNAKE_CASE_			).manual_seed(0			)
				lowerCAmelCase												=  pipe(
				    [prompt]      ,generator=SCREAMING_SNAKE_CASE_      ,output_type='''np'''      ,return_dict=SCREAMING_SNAKE_CASE_      ,num_inference_steps=2			)[0]

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

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

				lowerCAmelCase												=  np.array([0.66_93, 0.60_75, 0.49_59, 0.57_01, 0.55_83, 0.43_33, 0.61_71, 0.56_84, 0.49_88]			)

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






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





				"""simple docstring"""


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






			def 					UpperCAmelCase							(self  :	List[str]			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase												=  load_numpy(
				    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
				    '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''			)

				lowerCAmelCase												=  VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''			)
				lowerCAmelCase												=  pipeline.to(SCREAMING_SNAKE_CASE_			)
				pipeline.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_			)

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

				lowerCAmelCase												=  output.images[0]

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




"""simple docstring"""






def __magic_name__	(				_lowerCamelCase:		Tuple,					_lowerCamelCase:		Any			)					->						Optional[Any]:

	'''simple docstring'''




	lowerCAmelCase												=  [1]
	for i in range(2,					_lowerCamelCase			):
		factorials.append(factorials[-1] * i			)
	assert 0 <= k < factorials[-1] * n, "k out of bounds"

	lowerCAmelCase												=  []
	lowerCAmelCase												=  list(range(_lowerCamelCase			)			)

	# Find permutation
	while factorials:
		lowerCAmelCase												=  factorials.pop()
		lowerCAmelCase							,       lowerCAmelCase												=  divmod(_lowerCamelCase,					_lowerCamelCase			)
		permutation.append(elements[number]			)
		elements.remove(elements[number]			)
	permutation.append(elements[0]			)

	return permutation


if __name__ == "__main__":
					import doctest

					doctest.testmod()


