





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




					def 	lowerCAmelCase			(		self  )		->							Optional[int]:
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




					@property
					def 	lowerCAmelCase			(		self  )		->							Union[str, Any]:
										return 12




					@property
					def 	lowerCAmelCase			(		self  )		->							List[str]:
										return 12




					@property
					def 	lowerCAmelCase			(		self  )		->							int:
										return 32




					@property
					def 	lowerCAmelCase			(		self  )		->							Dict:
										torch.manual_seed(0  )
										_snake_case									=				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  )		->							Tuple:
										_snake_case									=				CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'  )
										return tokenizer




					@property
					def 	lowerCAmelCase			(		self  )		->							Optional[Any]:
										torch.manual_seed(0  )
										_snake_case									=				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(lowerCAmelCase_  )




					@property
					def 	lowerCAmelCase			(		self  )		->							int:
										torch.manual_seed(0  )

										_snake_case									=				12
										_snake_case									=				12

										_snake_case									=				{
										    '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									=				TransformeraDModel(**lowerCAmelCase_  )
										return model




					def 	lowerCAmelCase			(		self  )		->							Optional[int]:
										_snake_case									=				'cpu'

										_snake_case									=				self.dummy_vqvae
										_snake_case									=				self.dummy_text_encoder
										_snake_case									=				self.dummy_tokenizer
										_snake_case									=				self.dummy_transformer
										_snake_case									=				VQDiffusionScheduler(self.num_embed  )
										_snake_case									=				LearnedClassifierFreeSamplingEmbeddings(learnable=lowerCAmelCase_  )

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

										_snake_case									=				'teddy bear playing in the pool'

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

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

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

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

										_snake_case									=				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 	lowerCAmelCase			(		self  )		->							Dict:
										_snake_case									=				'cpu'

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

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

										_snake_case									=				'teddy bear playing in the pool'

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

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

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

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

										_snake_case									=				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 			UpperCamelCase_		(  unittest.TestCase					):




					def 	lowerCAmelCase			(		self  )		->							Dict:
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()






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

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

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

										_snake_case									=				output.images[0]

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










from typing import TYPE_CHECKING

from ...utils import _LazyModule


UpperCAmelCase_      							=     {"""tokenization_byt5""": ["""ByT5Tokenizer"""]}


if TYPE_CHECKING:
					from .tokenization_byta import ByTaTokenizer
else:
					import sys

					UpperCAmelCase_      							=     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



