






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

						"""simple docstring"""



						def        UpperCAmelCase__					(self:					Dict   )    ->     List[Any]:




											'''simple docstring'''



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



						@property
						def        UpperCAmelCase__					(self:					str   )    ->     str:




											'''simple docstring'''



											return 12



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




											'''simple docstring'''



											return 12



						@property
						def        UpperCAmelCase__					(self:					Tuple   )    ->     Tuple:




											'''simple docstring'''



											return 32



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




											'''simple docstring'''



											torch.manual_seed(0   )
											__a					:						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        UpperCAmelCase__					(self:					Any   )    ->     List[Any]:




											'''simple docstring'''



											__a					:						Optional[Any]      =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )
											return tokenizer



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




											'''simple docstring'''



											torch.manual_seed(0   )
											__a					:						int      =					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        UpperCAmelCase__					(self:					int   )    ->     str:




											'''simple docstring'''



											torch.manual_seed(0   )

											__a					:						int      =					12
											__a					:						Dict      =					12

											__a					:						Tuple      =					{
											    "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",
											}

											__a					:						List[str]      =					TransformeraDModel(**__UpperCAmelCase   )
											return model



						def        UpperCAmelCase__					(self:					int   )    ->     Dict:




											'''simple docstring'''



											__a					:						Optional[Any]      =					"cpu"

											__a					:						Dict      =					self.dummy_vqvae
											__a					:						int      =					self.dummy_text_encoder
											__a					:						int      =					self.dummy_tokenizer
											__a					:						str      =					self.dummy_transformer
											__a					:						List[str]      =					VQDiffusionScheduler(self.num_embed   )
											__a					:						int      =					LearnedClassifierFreeSamplingEmbeddings(learnable=__UpperCAmelCase   )

											__a					:						List[str]      =					VQDiffusionPipeline(
											    vqvae=__UpperCAmelCase ,					text_encoder=__UpperCAmelCase ,					tokenizer=__UpperCAmelCase ,					transformer=__UpperCAmelCase ,					scheduler=__UpperCAmelCase ,					learned_classifier_free_sampling_embeddings=__UpperCAmelCase ,					)
											__a					:						List[Any]      =					pipe.to(__UpperCAmelCase   )
											pipe.set_progress_bar_config(disable=__UpperCAmelCase   )

											__a					:						Union[str, Any]      =					"teddy bear playing in the pool"

											__a					:						str      =					torch.Generator(device=__UpperCAmelCase   ).manual_seed(0   )
											__a					:						int      =					pipe([prompt] ,					generator=__UpperCAmelCase ,					num_inference_steps=2 ,					output_type="np"   )
											__a					:						Dict      =					output.images

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

											__a					:						Any      =					image[0, -3:, -3:, -1]
											__a					:						Dict      =					image_from_tuple[0, -3:, -3:, -1]

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

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




											'''simple docstring'''



											__a					:						List[Any]      =					"cpu"

											__a					:						int      =					self.dummy_vqvae
											__a					:						int      =					self.dummy_text_encoder
											__a					:						Tuple      =					self.dummy_tokenizer
											__a					:						Optional[Any]      =					self.dummy_transformer
											__a					:						Union[str, Any]      =					VQDiffusionScheduler(self.num_embed   )
											__a					:						Dict      =					LearnedClassifierFreeSamplingEmbeddings(
											    learnable=__UpperCAmelCase ,					hidden_size=self.text_embedder_hidden_size ,					length=tokenizer.model_max_length   )

											__a					:						Optional[int]      =					VQDiffusionPipeline(
											    vqvae=__UpperCAmelCase ,					text_encoder=__UpperCAmelCase ,					tokenizer=__UpperCAmelCase ,					transformer=__UpperCAmelCase ,					scheduler=__UpperCAmelCase ,					learned_classifier_free_sampling_embeddings=__UpperCAmelCase ,					)
											__a					:						Union[str, Any]      =					pipe.to(__UpperCAmelCase   )
											pipe.set_progress_bar_config(disable=__UpperCAmelCase   )

											__a					:						List[Any]      =					"teddy bear playing in the pool"

											__a					:						Dict      =					torch.Generator(device=__UpperCAmelCase   ).manual_seed(0   )
											__a					:						List[str]      =					pipe([prompt] ,					generator=__UpperCAmelCase ,					num_inference_steps=2 ,					output_type="np"   )
											__a					:						int      =					output.images

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

											__a					:						List[Any]      =					image[0, -3:, -3:, -1]
											__a					:						Union[str, Any]      =					image_from_tuple[0, -3:, -3:, -1]

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

											__a					:						List[str]      =					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 							snake_case_       (    unittest.TestCase      ):

						"""simple docstring"""



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




											'''simple docstring'''



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



						def        UpperCAmelCase__					(self:					Union[str, Any]   )    ->     Optional[int]:




											'''simple docstring'''



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

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

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

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

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













from abc import ABC, abstractmethod
from argparse import ArgumentParser




class 							snake_case_       (    __UpperCamelCase      ):

						"""simple docstring"""



						@staticmethod
						@abstractmethod
						def        UpperCAmelCase__					(__UpperCAmelCase:					ArgumentParser   )    ->     Tuple:




											'''simple docstring'''



											raise NotImplementedError()



						@abstractmethod
						def        UpperCAmelCase__					(self:					List[str]   )    ->     List[str]:




											'''simple docstring'''



											raise NotImplementedError()





