

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


lowerCAmelCase__								=		False







class 	lowercase    (				unittest.TestCase							):


  """simple docstring"""




  def 		A__					(				self):
   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()


  @property
  def 		A__					(				self):
   return 12


  @property
  def 		A__					(				self):
   return 12


  @property
  def 		A__					(				self):
   return 32


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


  @property
  def 		A__					(				self):
   torch.manual_seed(0)
   _UpperCamelCase   :		Dict			=						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=10_00	,       )
   return CLIPTextModel(A_)


  @property
  def 		A__					(				self):
   torch.manual_seed(0)

   _UpperCamelCase   :		Optional[int]			=						12
   _UpperCamelCase   :		Optional[int]			=						12

   _UpperCamelCase   :		Dict			=						{
       "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",
   }

   _UpperCamelCase   :		Any			=						TransformeraDModel(**A_)
   return model


  def 		A__					(				self):
   _UpperCamelCase   :		Optional[Any]			=						"cpu"

   _UpperCamelCase   :		Union[str, Any]			=						self.dummy_vqvae
   _UpperCamelCase   :		Dict			=						self.dummy_text_encoder
   _UpperCamelCase   :		List[str]			=						self.dummy_tokenizer
   _UpperCamelCase   :		str			=						self.dummy_transformer
   _UpperCamelCase   :		Optional[int]			=						VQDiffusionScheduler(self.num_embed)
   _UpperCamelCase   :		Any			=						LearnedClassifierFreeSamplingEmbeddings(learnable=A_)

   _UpperCamelCase   :		Any			=						VQDiffusionPipeline(
       vqvae=A_	,       text_encoder=A_	,       tokenizer=A_	,       transformer=A_	,       scheduler=A_	,       learned_classifier_free_sampling_embeddings=A_	,       )
   _UpperCamelCase   :		Any			=						pipe.to(A_)
   pipe.set_progress_bar_config(disable=A_)

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

   _UpperCamelCase   :		int			=						torch.Generator(device=A_).manual_seed(0)
   _UpperCamelCase   :		List[str]			=						pipe([prompt]	,       generator=A_	,       num_inference_steps=2	,       output_type='np')
   _UpperCamelCase   :		List[str]			=						output.images

   _UpperCamelCase   :		Union[str, Any]			=						torch.Generator(device=A_).manual_seed(0)
   _UpperCamelCase   :		str			=						pipe(
       [prompt]	,       generator=A_	,       output_type='np'	,       return_dict=A_	,       num_inference_steps=2)[0]

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

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

   _UpperCamelCase   :		Tuple			=						np.array([0.6_5_5_1, 0.6_1_6_8, 0.5_0_0_8, 0.5_6_7_6, 0.5_6_5_9, 0.4_2_9_5, 0.6_0_7_3, 0.5_5_9_9, 0.4_9_9_2])

   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 		A__					(				self):
   _UpperCamelCase   :		str			=						"cpu"

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

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

   _UpperCamelCase   :		List[str]			=						"teddy bear playing in the pool"

   _UpperCamelCase   :		List[str]			=						torch.Generator(device=A_).manual_seed(0)
   _UpperCamelCase   :		List[str]			=						pipe([prompt]	,       generator=A_	,       num_inference_steps=2	,       output_type='np')
   _UpperCamelCase   :		Any			=						output.images

   _UpperCamelCase   :		List[Any]			=						torch.Generator(device=A_).manual_seed(0)
   _UpperCamelCase   :		int			=						pipe(
       [prompt]	,       generator=A_	,       output_type='np'	,       return_dict=A_	,       num_inference_steps=2)[0]

   _UpperCamelCase   :		Any			=						image[0, -3:, -3:, -1]
   _UpperCamelCase   :		Tuple			=						image_from_tuple[0, -3:, -3:, -1]

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

   _UpperCamelCase   :		Optional[Any]			=						np.array([0.6_6_9_3, 0.6_0_7_5, 0.4_9_5_9, 0.5_7_0_1, 0.5_5_8_3, 0.4_3_3_3, 0.6_1_7_1, 0.5_6_8_4, 0.4_9_8_8])

   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							):


  """simple docstring"""




  def 		A__					(				self):
   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()






  def 		A__					(				self):
   _UpperCamelCase   :		int			=						load_numpy(
       'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
       '/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy')

   _UpperCamelCase   :		List[str]			=						VQDiffusionPipeline.from_pretrained('microsoft/vq-diffusion-ithq')
   _UpperCamelCase   :		List[str]			=						pipeline.to(A_)
   pipeline.set_progress_bar_config(disable=A_)

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

   _UpperCamelCase   :		List[Any]			=						output.images[0]

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







from ...processing_utils import ProcessorMixin







class 	lowercase    (				_lowercase							):


		"""simple docstring"""




		a__					       =  ["image_processor", "feature_extractor"]
		a__					       =  "TvltImageProcessor"
		a__					       =  "TvltFeatureExtractor"


		def __init__(				self	,       __snake_case	,       __snake_case):
			super().__init__(image_processor=__snake_case	,       feature_extractor=__snake_case)

			_UpperCamelCase   :		List[str]			=						image_processor
			_UpperCamelCase   :		Dict			=						feature_extractor


		def __call__(				self	,       __snake_case=None	,       __snake_case=None	,       __snake_case=None	,       __snake_case=None	,       __snake_case=False	,       __snake_case=False	,       *__snake_case	,       **__snake_case	,       ):

			if images is None and audio is None:
				raise ValueError('You need to specify either an `images` or `audio` input to process.')

			_UpperCamelCase   :		Union[str, Any]			=						None
			if images is not None:
				_UpperCamelCase   :		Tuple			=						self.image_processor(__snake_case	,       mask_pixel=__snake_case	,       *__snake_case	,       **__snake_case)
			if images_mixed is not None:
				_UpperCamelCase   :		Union[str, Any]			=						self.image_processor(__snake_case	,       is_mixed=__snake_case	,       *__snake_case	,       **__snake_case)
			if audio is not None:
				_UpperCamelCase   :		Tuple			=						self.feature_extractor(
				    __snake_case	,       *__snake_case	,       sampling_rate=__snake_case	,       mask_audio=__snake_case	,       **__snake_case)

			_UpperCamelCase   :		Tuple			=						{}
			if audio is not None:
				output_dict.update(__snake_case)
			if images is not None:
				output_dict.update(__snake_case)
			if images_mixed_dict is not None:
				output_dict.update(__snake_case)
			return output_dict






		@property
		def 		A__					(				self):
			_UpperCamelCase   :		List[Any]			=						self.image_processor.model_input_names
			_UpperCamelCase   :		List[Any]			=						self.feature_extractor.model_input_names
			return list(dict.fromkeys(image_processor_input_names + feature_extractor_input_names))




