





"""simple docstring"""




import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()

class     __SCREAMING_SNAKE_CASE   (       _a				,						unittest.TestCase    ):
							snake_case       : Union[str, Any]         =    PriorTransformer
							snake_case       : List[Any]         =    """hidden_states"""






							@property
							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     4
														UpperCamelCase__			       =     8
														UpperCamelCase__			       =     7

														UpperCamelCase__			       =     floats_tensor((batch_size, embedding_dim)					).to(__lowerCAmelCase					)

														UpperCamelCase__			       =     floats_tensor((batch_size, embedding_dim)					).to(__lowerCAmelCase					)
														UpperCamelCase__			       =     floats_tensor((batch_size, num_embeddings, embedding_dim)					).to(__lowerCAmelCase					)

														return {
														    "hidden_states": hidden_states,
														    "timestep": 2,
														    "proj_embedding": proj_embedding,
														    "encoder_hidden_states": encoder_hidden_states,
														}






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase=0					):
														torch.manual_seed(__lowerCAmelCase					)
														UpperCamelCase__			       =     4
														UpperCamelCase__			       =     8
														UpperCamelCase__			       =     7

														UpperCamelCase__			       =     torch.randn((batch_size, embedding_dim)					).to(__lowerCAmelCase					)

														UpperCamelCase__			       =     torch.randn((batch_size, embedding_dim)					).to(__lowerCAmelCase					)
														UpperCamelCase__			       =     torch.randn((batch_size, num_embeddings, embedding_dim)					).to(__lowerCAmelCase					)

														return {
														    "hidden_states": hidden_states,
														    "timestep": 2,
														    "proj_embedding": proj_embedding,
														    "encoder_hidden_states": encoder_hidden_states,
														}






							@property
							def 			_lowerCamelCase						(  self					):
														return (4, 8)






							@property
							def 			_lowerCamelCase						(  self					):
														return (4, 8)






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     {
														    """num_attention_heads""": 2,
														    """attention_head_dim""": 4,
														    """num_layers""": 2,
														    """embedding_dim""": 8,
														    """num_embeddings""": 7,
														    """additional_embeddings""": 4,
														}
														UpperCamelCase__			       =     self.dummy_input
														return init_dict, inputs_dict






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__       ,       UpperCamelCase__			       =     PriorTransformer.from_pretrained(
														    """hf-internal-testing/prior-dummy"""					,	output_loading_info=__lowerCAmelCase					)
														self.assertIsNotNone(__lowerCAmelCase					)
														self.assertEqual(len(loading_info["""missing_keys"""]					)					,	0					)

														model.to(__lowerCAmelCase					)
														UpperCamelCase__			       =     model(**self.dummy_input					)[0]

														assert hidden_states is not None, "Make sure output is not None"






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__       ,       UpperCamelCase__			       =     self.prepare_init_args_and_inputs_for_common()

														UpperCamelCase__			       =     self.model_class(**__lowerCAmelCase					)
														UpperCamelCase__			       =     inspect.signature(model.forward					)
														# signature.parameters is an OrderedDict => so arg_names order is deterministic
														UpperCamelCase__			       =     [*signature.parameters.keys()]

														UpperCamelCase__			       =     ["""hidden_states""", """timestep"""]
														self.assertListEqual(arg_names[:2]					,	__lowerCAmelCase					)







							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     PriorTransformer.from_pretrained("""hf-internal-testing/prior-dummy"""					)
														UpperCamelCase__			       =     model.to(__lowerCAmelCase					)

														if hasattr(__lowerCAmelCase					,	"""set_default_attn_processor"""					):
																					model.set_default_attn_processor()

														UpperCamelCase__			       =     self.get_dummy_seed_input()

														with torch.no_grad():
																					UpperCamelCase__			       =     model(**__lowerCAmelCase					)[0]

														UpperCamelCase__			       =     output[0, :5].flatten().cpu()
														print(__lowerCAmelCase					)

														# Since the VAE Gaussian prior's generator is seeded on the appropriate device,
														# the expected output slices are not the same for CPU and GPU.
														UpperCamelCase__			       =     torch.tensor([-1.3436, -0.2870, 0.7538, 0.4368, -0.0239]					)
														self.assertTrue(torch_all_close(__lowerCAmelCase					,	__lowerCAmelCase					,	rtol=1E-2					)					)






@slow
class     __SCREAMING_SNAKE_CASE   (       unittest.TestCase    ):






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase=1					,	__lowerCAmelCase=768					,	__lowerCAmelCase=77					,	__lowerCAmelCase=0					):
														torch.manual_seed(__lowerCAmelCase					)
														UpperCamelCase__			       =     batch_size
														UpperCamelCase__			       =     embedding_dim
														UpperCamelCase__			       =     num_embeddings

														UpperCamelCase__			       =     torch.randn((batch_size, embedding_dim)					).to(__lowerCAmelCase					)

														UpperCamelCase__			       =     torch.randn((batch_size, embedding_dim)					).to(__lowerCAmelCase					)
														UpperCamelCase__			       =     torch.randn((batch_size, num_embeddings, embedding_dim)					).to(__lowerCAmelCase					)

														return {
														    "hidden_states": hidden_states,
														    "timestep": 2,
														    "proj_embedding": proj_embedding,
														    "encoder_hidden_states": encoder_hidden_states,
														}






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







							@parameterized.expand(
							    [
							        # fmt: off
							        [13, [-0.5861,  0.1283, -0.0931,  0.0882,  0.4476,  0.1329, -0.0498,  0.0640]],
							        [37, [-0.4913,  0.0110, -0.0483,  0.0541,  0.4954, -0.0170,  0.0354,  0.1651]],
							        # fmt: on
							    ]					)
							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     PriorTransformer.from_pretrained("""kandinsky-community/kandinsky-2-1-prior"""					,	subfolder="""prior"""					)
														model.to(__lowerCAmelCase					)
														UpperCamelCase__			       =     self.get_dummy_seed_input(seed=__lowerCAmelCase					)

														with torch.no_grad():
																					UpperCamelCase__			       =     model(**__lowerCAmelCase					)[0]

														assert list(sample.shape					) == [1, 768]

														UpperCamelCase__			       =     sample[0, :8].flatten().cpu()
														print(__lowerCAmelCase					)
														UpperCamelCase__			       =     torch.tensor(__lowerCAmelCase					)

														assert torch_all_close(__lowerCAmelCase					,	__lowerCAmelCase					,	atol=1E-3					)











import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import is_speech_available, is_vision_available
from transformers.testing_utils import require_torch


if is_vision_available():
					from transformers import TvltImageProcessor

if is_speech_available():
					from transformers import TvltFeatureExtractor

from transformers import TvltProcessor

@require_torch
class     __SCREAMING_SNAKE_CASE   (       unittest.TestCase    ):






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     """ZinengTang/tvlt-base"""
														UpperCamelCase__			       =     tempfile.mkdtemp()






							def 			_lowerCamelCase						(  self					,	**__lowerCAmelCase					):
														return TvltImageProcessor.from_pretrained(self.checkpoint					,	**__lowerCAmelCase					)






							def 			_lowerCamelCase						(  self					,	**__lowerCAmelCase					):
														return TvltFeatureExtractor.from_pretrained(self.checkpoint					,	**__lowerCAmelCase					)






							def 			_lowerCamelCase						(  self					):
														shutil.rmtree(self.tmpdirname					)






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.get_image_processor()
														UpperCamelCase__			       =     self.get_feature_extractor()

														UpperCamelCase__			       =     TvltProcessor(image_processor=__lowerCAmelCase					,	feature_extractor=__lowerCAmelCase					)
														processor.save_pretrained(self.tmpdirname					)
														UpperCamelCase__			       =     TvltProcessor.from_pretrained(self.tmpdirname					)

														self.assertIsInstance(processor.feature_extractor					,	__lowerCAmelCase					)
														self.assertIsInstance(processor.image_processor					,	__lowerCAmelCase					)






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.get_image_processor()
														UpperCamelCase__			       =     self.get_feature_extractor()

														UpperCamelCase__			       =     TvltProcessor(image_processor=__lowerCAmelCase					,	feature_extractor=__lowerCAmelCase					)

														UpperCamelCase__			       =     np.ones([12000]					)

														UpperCamelCase__			       =     feature_extractor(__lowerCAmelCase					,	return_tensors="""np"""					)
														UpperCamelCase__			       =     processor(audio=__lowerCAmelCase					,	return_tensors="""np"""					)

														for key in audio_dict.keys():
																					self.assertAlmostEqual(audio_dict[key].sum()					,	input_processor[key].sum()					,	delta=1E-2					)






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.get_image_processor()
														UpperCamelCase__			       =     self.get_feature_extractor()

														UpperCamelCase__			       =     TvltProcessor(image_processor=__lowerCAmelCase					,	feature_extractor=__lowerCAmelCase					)

														UpperCamelCase__			       =     np.ones([3, 224, 224]					)

														UpperCamelCase__			       =     image_processor(__lowerCAmelCase					,	return_tensors="""np"""					)
														UpperCamelCase__			       =     processor(images=__lowerCAmelCase					,	return_tensors="""np"""					)

														for key in image_dict.keys():
																					self.assertAlmostEqual(image_dict[key].sum()					,	input_processor[key].sum()					,	delta=1E-2					)






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.get_image_processor()
														UpperCamelCase__			       =     self.get_feature_extractor()

														UpperCamelCase__			       =     TvltProcessor(image_processor=__lowerCAmelCase					,	feature_extractor=__lowerCAmelCase					)

														UpperCamelCase__			       =     np.ones([12000]					)
														UpperCamelCase__			       =     np.ones([3, 224, 224]					)

														UpperCamelCase__			       =     processor(audio=__lowerCAmelCase					,	images=__lowerCAmelCase					)

														self.assertListEqual(list(inputs.keys()					)					,	["""audio_values""", """audio_mask""", """pixel_values""", """pixel_mask"""]					)

														# test if it raises when no input is passed
														with pytest.raises(__lowerCAmelCase					):
																					processor()







							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.get_image_processor()
														UpperCamelCase__			       =     self.get_feature_extractor()

														UpperCamelCase__			       =     TvltProcessor(image_processor=__lowerCAmelCase					,	feature_extractor=__lowerCAmelCase					)

														self.assertListEqual(
														    processor.model_input_names					,	image_processor.model_input_names + feature_extractor.model_input_names					,	msg="""`processor` and `image_processor`+`feature_extractor` model input names do not match"""					,	)




