import tempfile
import unittest

import numpy as np
import torch

from diffusers import DiffusionPipeline
from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor




class        _UpperCamelCase				(						unittest.TestCase			):


			"""simple docstring"""






			def 	_SCREAMING_SNAKE_CASE      (       self				,       lowerCAmelCase__ = False						)					-> Any:

							'''simple docstring'''


							__lowercase			      =							10

							if only_cross_attention:
											__lowercase			      =							12
							else:
											# when only cross attention is not set, the cross attention dim must be the same as the query dim
											__lowercase			      =							query_dim

							return {
							    "query_dim": query_dim,
							    "cross_attention_dim": cross_attention_dim,
							    "heads": 2,
							    "dim_head": 4,
							    "added_kv_proj_dim": 6,
							    "norm_num_groups": 1,
							    "only_cross_attention": only_cross_attention,
							    "processor": AttnAddedKVProcessor(),
							}





			def 	_SCREAMING_SNAKE_CASE      (       self				,       lowerCAmelCase__				,       lowerCAmelCase__						)					-> Optional[Any]:

							'''simple docstring'''


							__lowercase			      =							2

							__lowercase			      =							torch.rand(lowerCAmelCase__				,       lowerCAmelCase__				,       3				,       2						)
							__lowercase			      =							torch.rand(lowerCAmelCase__				,       4				,       lowerCAmelCase__						)
							__lowercase			      =							None

							return {
							    "hidden_states": hidden_states,
							    "encoder_hidden_states": encoder_hidden_states,
							    "attention_mask": attention_mask,
							}





			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> List[Any]:

							'''simple docstring'''


							torch.manual_seed(0						)

							__lowercase			      =							self.get_constructor_arguments(only_cross_attention=lowerCAmelCase__						)
							__lowercase			      =							Attention(**lowerCAmelCase__						)

							self.assertTrue(attn.to_k is not None						)
							self.assertTrue(attn.to_v is not None						)

							__lowercase			      =							self.get_forward_arguments(
							    query_dim=constructor_args['''query_dim''']				,       added_kv_proj_dim=constructor_args['''added_kv_proj_dim''']						)

							__lowercase			      =							attn(**lowerCAmelCase__						)

							# only self attention

							torch.manual_seed(0						)

							__lowercase			      =							self.get_constructor_arguments(only_cross_attention=lowerCAmelCase__						)
							__lowercase			      =							Attention(**lowerCAmelCase__						)

							self.assertTrue(attn.to_k is None						)
							self.assertTrue(attn.to_v is None						)

							__lowercase			      =							self.get_forward_arguments(
							    query_dim=constructor_args['''query_dim''']				,       added_kv_proj_dim=constructor_args['''added_kv_proj_dim''']						)

							__lowercase			      =							attn(**lowerCAmelCase__						)

							self.assertTrue((only_cross_attn_out != self_and_cross_attn_out).all()						)




class        _UpperCamelCase				(						unittest.TestCase			):


			"""simple docstring"""






			def 	_SCREAMING_SNAKE_CASE      (       self						)					-> Optional[Any]:

							'''simple docstring'''


							__lowercase			      =							DiffusionPipeline.from_pretrained('''hf-internal-testing/tiny-stable-diffusion-pipe'''				,       safety_checker=lowerCAmelCase__						)

							__lowercase			      =							pipe(
							    '''foo'''				,       num_inference_steps=2				,       generator=torch.Generator('''cpu'''						).manual_seed(0						)				,       output_type='''np'''				,       ).images

							# the initial conversion succeeds
							__lowercase			      =							DiffusionPipeline.from_pretrained(
							    '''hf-internal-testing/tiny-stable-diffusion-pipe'''				,       device_map='''sequential'''				,       safety_checker=lowerCAmelCase__						)

							__lowercase			      =							pipe(
							    '''foo'''				,       num_inference_steps=2				,       generator=torch.Generator('''cpu'''						).manual_seed(0						)				,       output_type='''np'''				,       ).images

							with tempfile.TemporaryDirectory() as tmpdir:
											# save the converted model
											pipe.save_pretrained(lowerCAmelCase__						)

											# can also load the converted weights
											__lowercase			      =							DiffusionPipeline.from_pretrained(lowerCAmelCase__				,       device_map='''sequential'''				,       safety_checker=lowerCAmelCase__						)

							__lowercase			      =							pipe(
							    '''foo'''				,       num_inference_steps=2				,       generator=torch.Generator('''cpu'''						).manual_seed(0						)				,       output_type='''np'''				,       ).images

							self.assertTrue(np.allclose(lowerCAmelCase__				,       lowerCAmelCase__				,       atol=1E-5						)						)
							self.assertTrue(np.allclose(lowerCAmelCase__				,       lowerCAmelCase__				,       atol=1E-5						)						)
import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


__a  :					int          =  5_0_0_0_0_0

__a   ,						__a  :					Any          =  os.path.split(__file__)
__a  :					Dict          =  os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))





@get_duration
def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							dataset.select(range(0							,      len(lowercase							)							,      2							)							)





@get_duration
def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							dataset.sort('''numbers'''							)





@get_duration
def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							dataset.shuffle()





@get_duration
def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							dataset.train_test_split(0.1							)





@get_duration
def     UpperCAmelCase						(   lowercase							,      lowercase=10							):







				"""simple docstring"""



				for shard_id in range(lowercase							):
								__lowercase			      =							dataset.shard(lowercase							,      lowercase							)





def     UpperCAmelCase						(   ):







				"""simple docstring"""



				__lowercase			      =							{'''num examples''': SPEED_TEST_N_EXAMPLES}
				__lowercase			      =							(select, sort, shuffle, train_test_split, shard)
				with tempfile.TemporaryDirectory() as tmp_dir:
								print('''generating dataset'''							)
								__lowercase			      =							datasets.Features({'''text''': datasets.Value('''string'''							), '''numbers''': datasets.Value('''float32'''							)}							)
								__lowercase			      =							generate_example_dataset(
								    os.path.join(lowercase							,      '''dataset.arrow'''							)							,      lowercase							,      num_examples=lowercase							)
								print('''Functions'''							)
								for func in functions:
												print(func.__name__							)
												__lowercase			      =							func(lowercase							)

				with open(lowercase							,      '''wb'''							) as f:
								f.write(json.dumps(lowercase							).encode('''utf-8'''							)							)


if __name__ == "__main__":  # useful to run the profiler
		benchmark_indices_mapping()