



import random
import unittest

import torch

from diffusers import IFSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin


@skip_mps
class    _lowercase (a_   , a_   , unittest.TestCase   ):
				'''simple docstring'''







				lowercase__     =    IFSuperResolutionPipeline
				lowercase__     =    TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""width""", """height"""}
				lowercase__     =    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
				lowercase__     =    PipelineTesterMixin.required_optional_params - {"""latents"""}

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									return self._get_superresolution_dummy_components()

				def    _lowerCamelCase       (			self  ,			snake_case__  ,			snake_case__=0     ):

									'''simple docstring'''


									if str(snake_case__     ).startswith("mps"     ):
														UpperCamelCase_			= torch.manual_seed(snake_case__     )
									else:
														UpperCamelCase_			= torch.Generator(device=snake_case__     ).manual_seed(snake_case__     )

									UpperCamelCase_			= floats_tensor((1, 3, 32, 32)  ,			rng=random.Random(snake_case__     )     ).to(snake_case__     )

									UpperCamelCase_			= {
									    "prompt": "A painting of a squirrel eating a burger",
									    "image": image,
									    "generator": generator,
									    "num_inference_steps": 2,
									    "output_type": "numpy",
									}

									return inputs

				@unittest.skipIf(
				    torch_device != "cuda" or not is_xformers_available()  ,			reason="XFormers attention is only available with CUDA and `xformers` installed"  ,			)
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3     )

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									self._test_save_load_optional_components()

				@unittest.skipIf(torch_device != "cuda"  ,			reason="float16 requires CUDA"     )
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									super().test_save_load_floataa(expected_max_diff=1e-1     )

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									self._test_attention_slicing_forward_pass(expected_max_diff=1e-2     )

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									self._test_save_load_local()

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									self._test_inference_batch_single_identical(
									    expected_max_diff=1e-2  ,			)





import json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


UpperCAmelCase :				Any							=1000_0000_0000
UpperCAmelCase :				Any							=1_0000

UpperCAmelCase			,			UpperCAmelCase :				Tuple							=os.path.split(__file__)
UpperCAmelCase :				Tuple							=os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))



def  _lowerCAmelCase     (_lowerCAmelCase ,						_lowerCAmelCase):
					UpperCamelCase_			= pa.Table.from_pydict({"col": [0] * chunk_size})
					UpperCamelCase_			= pa.concat_tables([table] * (num_examples // chunk_size))
					return datasets.Dataset(_lowerCAmelCase ,						fingerprint="table_100B")


@dataclass
class    _lowercase :
				'''simple docstring'''







				lowercase__     =    42
				lowercase__     =    42
				lowercase__     =    42
				lowercase__     =    42

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= np.random.default_rng(self.seed     )
									UpperCamelCase_			= rng.integers(low=self.low  ,			high=self.high  ,			size=self.size     ).tolist()

				def __iter__(			self     ):

									'''simple docstring'''


									return iter(self._sampled_values     )

				def __len__(			self     ):

									'''simple docstring'''


									return self.size




@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[0]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[-1]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[range(len(_lowerCAmelCase) // 2 ,						len(_lowerCAmelCase) // 2 + 10_24)]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[RandIter(0 ,						len(_lowerCAmelCase) ,						10_24 ,						seed=42)]



def  _lowerCAmelCase     ():
					UpperCamelCase_			= {"num examples": SPEED_TEST_N_EXAMPLES}
					UpperCamelCase_			= (get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
					print("generating dataset")
					UpperCamelCase_			= generate_100B_dataset(num_examples=_lowerCAmelCase ,						chunk_size=_lowerCAmelCase)
					print("Functions")
					for func in functions:
										print(func.__name__)
										UpperCamelCase_			= func(_lowerCAmelCase)

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


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