





"""simple docstring"""


import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device





def 							a__	(     _UpperCAmelCase		)			->     Any:

 '''simple docstring'''






 snake_case__										= {}
 for name in unet.attn_processors.keys():
  snake_case__										= None if name.endswith("attn1.processor"		) else unet.config.cross_attention_dim
  if name.startswith("mid_block"		):
   snake_case__										= unet.config.block_out_channels[-1]
  elif name.startswith("up_blocks"		):
   snake_case__										= int(name[len("up_blocks."		)]		)
   snake_case__										= list(reversed(unet.config.block_out_channels		)		)[block_id]
  elif name.startswith("down_blocks"		):
   snake_case__										= int(name[len("down_blocks."		)]		)
   snake_case__										= unet.config.block_out_channels[block_id]
  snake_case__										= (
      LoRAAttnProcessor2_0 if hasattr(_UpperCAmelCase		,   "scaled_dot_product_attention"		) else LoRAAttnProcessor
  )
  snake_case__										= lora_attn_processor_class(
      hidden_size=_UpperCAmelCase		,   cross_attention_dim=_UpperCAmelCase		)
 snake_case__										= AttnProcsLayers(_UpperCAmelCase		)
 return lora_attn_procs, unet_lora_layers





def 							a__	(     _UpperCAmelCase		)			->     List[Any]:

 '''simple docstring'''






 snake_case__										= {}
 snake_case__										= (
     LoRAAttnProcessor2_0 if hasattr(_UpperCAmelCase		,   "scaled_dot_product_attention"		) else LoRAAttnProcessor
 )
 for name, module in text_encoder.named_modules():
  if name.endswith(_UpperCAmelCase		):
   snake_case__										= lora_attn_processor_class(
       hidden_size=module.out_proj.out_features		,   cross_attention_dim=_UpperCAmelCase		)
 return text_lora_attn_procs





def 							a__	(     _UpperCAmelCase		)			->     Tuple:

 '''simple docstring'''






 snake_case__										= create_text_encoder_lora_attn_procs(_UpperCAmelCase		)
 snake_case__										= AttnProcsLayers(_UpperCAmelCase		)
 return text_encoder_lora_layers





def 							a__	(     _UpperCAmelCase		,   _UpperCAmelCase=False		)			->     Optional[int]:

 '''simple docstring'''






 for _, attn_proc in text_lora_attn_procs.items():
  # set up.weights
  for layer_name, layer_module in attn_proc.named_modules():
   if layer_name.endswith("_lora"		):
    snake_case__										= (
        torch.randn_like(layer_module.up.weight		)
        if randn_weight
        else torch.zeros_like(layer_module.up.weight		)
    )
    snake_case__										= torch.nn.Parameter(_UpperCAmelCase		)



class    lowercase_ (   unittest.TestCase							):

       def 						_SCREAMING_SNAKE_CASE(	self   ):






        '''simple docstring'''

        torch.manual_seed(0   )
        snake_case__										= UNetaDConditionModel(
            block_out_channels=(32, 64)      ,    layers_per_block=2      ,    sample_size=32      ,    in_channels=4      ,    out_channels=4      ,    down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")      ,    up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")      ,    cross_attention_dim=32      ,    )
        snake_case__										= DDIMScheduler(
            beta_start=0.00085      ,    beta_end=0.012      ,    beta_schedule="scaled_linear"      ,    clip_sample=__SCREAMING_SNAKE_CASE      ,    set_alpha_to_one=__SCREAMING_SNAKE_CASE      ,    steps_offset=1      ,    )
        torch.manual_seed(0   )
        snake_case__										= AutoencoderKL(
            block_out_channels=[32, 64]      ,    in_channels=3      ,    out_channels=3      ,    down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]      ,    up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]      ,    latent_channels=4      ,    )
        snake_case__										= CLIPTextConfig(
            bos_token_id=0      ,    eos_token_id=2      ,    hidden_size=32      ,    intermediate_size=37      ,    layer_norm_eps=1e-05      ,    num_attention_heads=4      ,    num_hidden_layers=5      ,    pad_token_id=1      ,    vocab_size=1000      ,    )
        snake_case__										= CLIPTextModel(__SCREAMING_SNAKE_CASE   )
        snake_case__										= CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )

        snake_case__							,      snake_case__										= create_unet_lora_layers(__SCREAMING_SNAKE_CASE   )
        snake_case__										= create_text_encoder_lora_layers(__SCREAMING_SNAKE_CASE   )

        snake_case__										= {
            "unet": unet,
            "scheduler": scheduler,
            "vae": vae,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "safety_checker": None,
            "feature_extractor": None,
        }
        snake_case__										= {
            "unet_lora_layers": unet_lora_layers,
            "text_encoder_lora_layers": text_encoder_lora_layers,
            "unet_lora_attn_procs": unet_lora_attn_procs,
        }
        return pipeline_components, lora_components

       def 						_SCREAMING_SNAKE_CASE(	self   ):






        '''simple docstring'''

        snake_case__										= 1
        snake_case__										= 10
        snake_case__										= 4
        snake_case__										= (32, 32)

        snake_case__										= torch.manual_seed(0   )
        snake_case__										= floats_tensor((batch_size, num_channels) + sizes   )
        snake_case__										= torch.randint(1      ,    __SCREAMING_SNAKE_CASE      ,    size=(batch_size, sequence_length)      ,    generator=__SCREAMING_SNAKE_CASE   )

        snake_case__										= {
            "prompt": "A painting of a squirrel eating a burger",
            "generator": generator,
            "num_inference_steps": 2,
            "guidance_scale": 6.0,
            "output_type": "numpy",
        }

        return noise, input_ids, pipeline_inputs

        # copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb

       def 						_SCREAMING_SNAKE_CASE(	self   ):






        '''simple docstring'''

        snake_case__										= 77

        snake_case__										= torch.randint(2      ,    56      ,    size=(1, max_seq_length)      ,    generator=torch.manual_seed(0   )   )

        snake_case__										= {}
        snake_case__										= inputs
        return prepared_inputs

       def 						_SCREAMING_SNAKE_CASE(	self      ,    __SCREAMING_SNAKE_CASE   ):






        '''simple docstring'''

        snake_case__							,      snake_case__										= self.get_dummy_components()
        LoraLoaderMixin.save_lora_weights(
            save_directory=__SCREAMING_SNAKE_CASE      ,    unet_lora_layers=lora_components["unet_lora_layers"]      ,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]      ,    )
        self.assertTrue(os.path.isfile(os.path.join(__SCREAMING_SNAKE_CASE      ,    "pytorch_lora_weights.bin"   )   )   )

       def 						_SCREAMING_SNAKE_CASE(	self   ):






        '''simple docstring'''

        snake_case__							,      snake_case__										= self.get_dummy_components()
        snake_case__										= StableDiffusionPipeline(**__SCREAMING_SNAKE_CASE   )
        snake_case__										= sd_pipe.to(__SCREAMING_SNAKE_CASE   )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE   )

        snake_case__							,      snake_case__							,      snake_case__										= self.get_dummy_inputs()

        snake_case__										= sd_pipe(**__SCREAMING_SNAKE_CASE   ).images
        snake_case__										= original_images[0, -3:, -3:, -1]

        with tempfile.TemporaryDirectory() as tmpdirname:
         LoraLoaderMixin.save_lora_weights(
             save_directory=__SCREAMING_SNAKE_CASE      ,    unet_lora_layers=lora_components["unet_lora_layers"]      ,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]      ,    )
         self.assertTrue(os.path.isfile(os.path.join(__SCREAMING_SNAKE_CASE      ,    "pytorch_lora_weights.bin"   )   )   )
         sd_pipe.load_lora_weights(__SCREAMING_SNAKE_CASE   )

        snake_case__										= sd_pipe(**__SCREAMING_SNAKE_CASE   ).images
        snake_case__										= lora_images[0, -3:, -3:, -1]

        # Outputs shouldn't match.
        self.assertFalse(torch.allclose(torch.from_numpy(__SCREAMING_SNAKE_CASE   )      ,    torch.from_numpy(__SCREAMING_SNAKE_CASE   )   )   )

       def 						_SCREAMING_SNAKE_CASE(	self   ):






        '''simple docstring'''

        snake_case__							,      snake_case__										= self.get_dummy_components()
        snake_case__										= StableDiffusionPipeline(**__SCREAMING_SNAKE_CASE   )
        snake_case__										= sd_pipe.to(__SCREAMING_SNAKE_CASE   )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE   )

        snake_case__							,      snake_case__							,      snake_case__										= self.get_dummy_inputs()

        snake_case__										= sd_pipe(**__SCREAMING_SNAKE_CASE   ).images
        snake_case__										= original_images[0, -3:, -3:, -1]

        with tempfile.TemporaryDirectory() as tmpdirname:
         LoraLoaderMixin.save_lora_weights(
             save_directory=__SCREAMING_SNAKE_CASE      ,    unet_lora_layers=lora_components["unet_lora_layers"]      ,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]      ,    safe_serialization=__SCREAMING_SNAKE_CASE      ,    )
         self.assertTrue(os.path.isfile(os.path.join(__SCREAMING_SNAKE_CASE      ,    "pytorch_lora_weights.safetensors"   )   )   )
         sd_pipe.load_lora_weights(__SCREAMING_SNAKE_CASE   )

        snake_case__										= sd_pipe(**__SCREAMING_SNAKE_CASE   ).images
        snake_case__										= lora_images[0, -3:, -3:, -1]

        # Outputs shouldn't match.
        self.assertFalse(torch.allclose(torch.from_numpy(__SCREAMING_SNAKE_CASE   )      ,    torch.from_numpy(__SCREAMING_SNAKE_CASE   )   )   )

       def 						_SCREAMING_SNAKE_CASE(	self   ):






        '''simple docstring'''

        snake_case__							,      snake_case__										= self.get_dummy_components()
        snake_case__										= lora_components["unet_lora_attn_procs"]
        snake_case__										= StableDiffusionPipeline(**__SCREAMING_SNAKE_CASE   )
        snake_case__										= sd_pipe.to(__SCREAMING_SNAKE_CASE   )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE   )

        snake_case__							,      snake_case__							,      snake_case__										= self.get_dummy_inputs()

        snake_case__										= sd_pipe(**__SCREAMING_SNAKE_CASE   ).images
        snake_case__										= original_images[0, -3:, -3:, -1]

        with tempfile.TemporaryDirectory() as tmpdirname:
         snake_case__										= sd_pipe.unet
         unet.set_attn_processor(__SCREAMING_SNAKE_CASE   )
         unet.save_attn_procs(__SCREAMING_SNAKE_CASE   )
         self.assertTrue(os.path.isfile(os.path.join(__SCREAMING_SNAKE_CASE      ,    "pytorch_lora_weights.bin"   )   )   )
         sd_pipe.load_lora_weights(__SCREAMING_SNAKE_CASE   )

        snake_case__										= sd_pipe(**__SCREAMING_SNAKE_CASE   ).images
        snake_case__										= lora_images[0, -3:, -3:, -1]

        # Outputs shouldn't match.
        self.assertFalse(torch.allclose(torch.from_numpy(__SCREAMING_SNAKE_CASE   )      ,    torch.from_numpy(__SCREAMING_SNAKE_CASE   )   )   )

       def 						_SCREAMING_SNAKE_CASE(	self   ):






        '''simple docstring'''

        snake_case__							,      snake_case__										= self.get_dummy_components()
        snake_case__										= StableDiffusionPipeline(**__SCREAMING_SNAKE_CASE   )

        snake_case__										= self.get_dummy_tokens()

        # inference without lora
        snake_case__										= pipe.text_encoder(**__SCREAMING_SNAKE_CASE   )[0]
        assert outputs_without_lora.shape == (1, 77, 32)

        # create lora_attn_procs with zeroed out up.weights
        snake_case__										= create_text_encoder_lora_attn_procs(pipe.text_encoder   )
        set_lora_up_weights(__SCREAMING_SNAKE_CASE      ,    randn_weight=__SCREAMING_SNAKE_CASE   )

        # monkey patch
        pipe._modify_text_encoder(__SCREAMING_SNAKE_CASE   )

        # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
        del text_attn_procs
        gc.collect()

        # inference with lora
        snake_case__										= pipe.text_encoder(**__SCREAMING_SNAKE_CASE   )[0]
        assert outputs_with_lora.shape == (1, 77, 32)

        assert torch.allclose(
            __SCREAMING_SNAKE_CASE      ,    __SCREAMING_SNAKE_CASE   ), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

        # create lora_attn_procs with randn up.weights
        snake_case__										= create_text_encoder_lora_attn_procs(pipe.text_encoder   )
        set_lora_up_weights(__SCREAMING_SNAKE_CASE      ,    randn_weight=__SCREAMING_SNAKE_CASE   )

        # monkey patch
        pipe._modify_text_encoder(__SCREAMING_SNAKE_CASE   )

        # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
        del text_attn_procs
        gc.collect()

        # inference with lora
        snake_case__										= pipe.text_encoder(**__SCREAMING_SNAKE_CASE   )[0]
        assert outputs_with_lora.shape == (1, 77, 32)

        assert not torch.allclose(
            __SCREAMING_SNAKE_CASE      ,    __SCREAMING_SNAKE_CASE   ), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"

       def 						_SCREAMING_SNAKE_CASE(	self   ):






        '''simple docstring'''

        snake_case__							,      snake_case__										= self.get_dummy_components()
        snake_case__										= StableDiffusionPipeline(**__SCREAMING_SNAKE_CASE   )

        snake_case__										= self.get_dummy_tokens()

        # inference without lora
        snake_case__										= pipe.text_encoder(**__SCREAMING_SNAKE_CASE   )[0]
        assert outputs_without_lora.shape == (1, 77, 32)

        # create lora_attn_procs with randn up.weights
        snake_case__										= create_text_encoder_lora_attn_procs(pipe.text_encoder   )
        set_lora_up_weights(__SCREAMING_SNAKE_CASE      ,    randn_weight=__SCREAMING_SNAKE_CASE   )

        # monkey patch
        pipe._modify_text_encoder(__SCREAMING_SNAKE_CASE   )

        # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
        del text_attn_procs
        gc.collect()

        # inference with lora
        snake_case__										= pipe.text_encoder(**__SCREAMING_SNAKE_CASE   )[0]
        assert outputs_with_lora.shape == (1, 77, 32)

        assert not torch.allclose(
            __SCREAMING_SNAKE_CASE      ,    __SCREAMING_SNAKE_CASE   ), "lora outputs should be different to without lora outputs"

        # remove monkey patch
        pipe._remove_text_encoder_monkey_patch()

        # inference with removed lora
        snake_case__										= pipe.text_encoder(**__SCREAMING_SNAKE_CASE   )[0]
        assert outputs_without_lora_removed.shape == (1, 77, 32)

        assert torch.allclose(
            __SCREAMING_SNAKE_CASE      ,    __SCREAMING_SNAKE_CASE   ), "remove lora monkey patch should restore the original outputs"

       def 						_SCREAMING_SNAKE_CASE(	self   ):






        '''simple docstring'''

        snake_case__							,      snake_case__										= self.get_dummy_components()
        snake_case__										= StableDiffusionPipeline(**__SCREAMING_SNAKE_CASE   )
        snake_case__										= sd_pipe.to(__SCREAMING_SNAKE_CASE   )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE   )

        snake_case__							,      snake_case__							,      snake_case__										= self.get_dummy_inputs()

        with tempfile.TemporaryDirectory() as tmpdirname:
         LoraLoaderMixin.save_lora_weights(
             save_directory=__SCREAMING_SNAKE_CASE      ,    unet_lora_layers=lora_components["unet_lora_layers"]      ,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]      ,    )
         self.assertTrue(os.path.isfile(os.path.join(__SCREAMING_SNAKE_CASE      ,    "pytorch_lora_weights.bin"   )   )   )
         sd_pipe.load_lora_weights(__SCREAMING_SNAKE_CASE   )

        snake_case__										= sd_pipe(**__SCREAMING_SNAKE_CASE   ).images
        snake_case__										= lora_images[0, -3:, -3:, -1]

        snake_case__										= sd_pipe(**__SCREAMING_SNAKE_CASE      ,    cross_attention_kwargs={"scale": 0.5}   ).images
        snake_case__										= lora_images_with_scale[0, -3:, -3:, -1]

        # Outputs shouldn't match.
        self.assertFalse(
            torch.allclose(torch.from_numpy(__SCREAMING_SNAKE_CASE   )      ,    torch.from_numpy(__SCREAMING_SNAKE_CASE   )   )   )

       def 						_SCREAMING_SNAKE_CASE(	self   ):






        '''simple docstring'''

        with tempfile.TemporaryDirectory() as tmpdirname:
         self.create_lora_weight_file(__SCREAMING_SNAKE_CASE   )

         snake_case__							,      snake_case__										= self.get_dummy_components()
         snake_case__										= StableDiffusionPipeline(**__SCREAMING_SNAKE_CASE   )
         snake_case__										= sd_pipe.to(__SCREAMING_SNAKE_CASE   )
         sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE   )

         # check if vanilla attention processors are used
         for _, module in sd_pipe.unet.named_modules():
          if isinstance(__SCREAMING_SNAKE_CASE      ,    __SCREAMING_SNAKE_CASE   ):
           self.assertIsInstance(module.processor      ,    (AttnProcessor, AttnProcessor2_0)   )

            # load LoRA weight file
         sd_pipe.load_lora_weights(__SCREAMING_SNAKE_CASE   )

         # check if lora attention processors are used
         for _, module in sd_pipe.unet.named_modules():
          if isinstance(__SCREAMING_SNAKE_CASE      ,    __SCREAMING_SNAKE_CASE   ):
           snake_case__										= (
               LoRAAttnProcessor2_0 if hasattr(__SCREAMING_SNAKE_CASE      ,    "scaled_dot_product_attention"   ) else LoRAAttnProcessor
           )
           self.assertIsInstance(module.processor      ,    __SCREAMING_SNAKE_CASE   )

       @unittest.skipIf(torch_device != "cuda"      ,    "This test is supposed to run on GPU"   )
       def 						_SCREAMING_SNAKE_CASE(	self   ):






        '''simple docstring'''

        with tempfile.TemporaryDirectory() as tmpdirname:
         self.create_lora_weight_file(__SCREAMING_SNAKE_CASE   )

         snake_case__							,      snake_case__										= self.get_dummy_components()
         snake_case__										= StableDiffusionPipeline(**__SCREAMING_SNAKE_CASE   )
         snake_case__										= sd_pipe.to(__SCREAMING_SNAKE_CASE   )
         sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE   )

         # enable XFormers
         sd_pipe.enable_xformers_memory_efficient_attention()

         # check if xFormers attention processors are used
         for _, module in sd_pipe.unet.named_modules():
          if isinstance(__SCREAMING_SNAKE_CASE      ,    __SCREAMING_SNAKE_CASE   ):
           self.assertIsInstance(module.processor      ,    __SCREAMING_SNAKE_CASE   )

            # load LoRA weight file
         sd_pipe.load_lora_weights(__SCREAMING_SNAKE_CASE   )

         # check if lora attention processors are used
         for _, module in sd_pipe.unet.named_modules():
          if isinstance(__SCREAMING_SNAKE_CASE      ,    __SCREAMING_SNAKE_CASE   ):
           self.assertIsInstance(module.processor      ,    __SCREAMING_SNAKE_CASE   )

       @unittest.skipIf(torch_device != "cuda"      ,    "This test is supposed to run on GPU"   )
       def 						_SCREAMING_SNAKE_CASE(	self   ):






        '''simple docstring'''

        snake_case__							,      snake_case__										= self.get_dummy_components()
        snake_case__										= StableDiffusionPipeline(**__SCREAMING_SNAKE_CASE   )
        snake_case__										= sd_pipe.to(__SCREAMING_SNAKE_CASE   )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE   )

        snake_case__							,      snake_case__							,      snake_case__										= self.get_dummy_inputs()

        # enable XFormers
        sd_pipe.enable_xformers_memory_efficient_attention()

        snake_case__										= sd_pipe(**__SCREAMING_SNAKE_CASE   ).images
        snake_case__										= original_images[0, -3:, -3:, -1]

        with tempfile.TemporaryDirectory() as tmpdirname:
         LoraLoaderMixin.save_lora_weights(
             save_directory=__SCREAMING_SNAKE_CASE      ,    unet_lora_layers=lora_components["unet_lora_layers"]      ,    text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]      ,    )
         self.assertTrue(os.path.isfile(os.path.join(__SCREAMING_SNAKE_CASE      ,    "pytorch_lora_weights.bin"   )   )   )
         sd_pipe.load_lora_weights(__SCREAMING_SNAKE_CASE   )

        snake_case__										= sd_pipe(**__SCREAMING_SNAKE_CASE   ).images
        snake_case__										= lora_images[0, -3:, -3:, -1]

        # Outputs shouldn't match.
        self.assertFalse(torch.allclose(torch.from_numpy(__SCREAMING_SNAKE_CASE   )      ,    torch.from_numpy(__SCREAMING_SNAKE_CASE   )   )   )








'''simple docstring'''







import json
import os
import tempfile

import datasets
from datasets.arrow_writer import ArrowWriter
from datasets.features import ArrayaD
from utils import generate_examples, get_duration


_a					: Dict      =    (3_0, 4_8_7)
_a					: int      =    (3_6, 1_0_2_4)
_a					: str      =    (1_0_0, 1_0_0)
_a					: Tuple      =    1_0_0

_a					: List[str]      =    datasets.Features(
    {"""text""": ArrayaD(SHAPE_TEST_1, dtype="""float32"""), """image""": ArrayaD(SHAPE_TEST_2, dtype="""float32""")}
)

_a    ,_a					: Optional[int]      =    os.path.split(__file__)
_a					: List[Any]      =    os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))





@get_duration
def    _lowerCAmelCase						(						lowercase					,	lowercase					,	lowercase	)					->					Optional[Any]:
			with ArrowWriter(features=lowercase					,	path=os.path.join(lowercase					,	"""beta.arrow"""	)	) as writer:
						for key, record in dummy_data:
									__lowerCAmelCase											=				my_features.encode_example(lowercase	)
									writer.write(lowercase	)
						__lowerCAmelCase						,       __lowerCAmelCase											=				writer.finalize()





@get_duration
def    _lowerCAmelCase						(						lowercase					,	lowercase	)					->					List[str]:
			__lowerCAmelCase											=				datasets.Dataset.from_file(
			    filename=os.path.join(lowercase					,	"""beta.arrow"""	)					,	info=datasets.DatasetInfo(features=lowercase	)	)
			for _ in dataset:
						pass





@get_duration
def    _lowerCAmelCase						(						lowercase					,	lowercase	)					->					Any:
			__lowerCAmelCase											=				datasets.Dataset.from_file(
			    filename=os.path.join(lowercase					,	"""beta.arrow"""	)					,	info=datasets.DatasetInfo(features=lowercase	)	)
			dataset.set_format("""numpy"""	)
			for _ in dataset:
						pass





@get_duration
def    _lowerCAmelCase						(						lowercase					,	lowercase	)					->					List[str]:
			__lowerCAmelCase											=				10
			__lowerCAmelCase											=				datasets.Dataset.from_file(
			    filename=os.path.join(lowercase					,	"""beta.arrow"""	)					,	info=datasets.DatasetInfo(features=lowercase	)	)
			for i in range(0					,	len(lowercase	)					,	lowercase	):
						__lowerCAmelCase											=				dataset[i : i + batch_size]





@get_duration
def    _lowerCAmelCase						(						lowercase					,	lowercase	)					->					List[str]:
			__lowerCAmelCase											=				10
			__lowerCAmelCase											=				datasets.Dataset.from_file(
			    filename=os.path.join(lowercase					,	"""beta.arrow"""	)					,	info=datasets.DatasetInfo(features=lowercase	)	)
			dataset.set_format("""numpy"""	)
			for i in range(0					,	len(lowercase	)					,	lowercase	):
						__lowerCAmelCase											=				dataset[i : i + batch_size]





@get_duration
def    _lowerCAmelCase						(						lowercase					,	lowercase	)					->					Union[str, Any]:
			__lowerCAmelCase											=				datasets.Dataset.from_file(
			    filename=os.path.join(lowercase					,	"""beta.arrow"""	)					,	info=datasets.DatasetInfo(features=lowercase	)	)
			for col in feats:
						__lowerCAmelCase											=				dataset[col]





@get_duration
def    _lowerCAmelCase						(						lowercase					,	lowercase	)					->					Optional[int]:
			__lowerCAmelCase											=				datasets.Dataset.from_file(
			    filename=os.path.join(lowercase					,	"""beta.arrow"""	)					,	info=datasets.DatasetInfo(features=lowercase	)	)
			dataset.set_format("""numpy"""	)
			for col in feats:
						__lowerCAmelCase											=				dataset[col]





def    _lowerCAmelCase						(						)					->					Optional[Any]:
			__lowerCAmelCase											=				{}
			__lowerCAmelCase											=				(
			    read_unformated,
			    read_formatted_as_numpy,
			    read_batch_unformated,
			    read_batch_formatted_as_numpy,
			    read_col_unformated,
			    read_col_formatted_as_numpy,
			)
			with tempfile.TemporaryDirectory() as tmp_dir:
						__lowerCAmelCase											=				datasets.Features({"""image""": ArrayaD(lowercase					,	dtype="""float32"""	)}	)
						__lowerCAmelCase											=				generate_examples(features=lowercase					,	num_examples=lowercase	)
						__lowerCAmelCase											=				write(lowercase					,	lowercase					,	lowercase	)
						for read_func in read_functions:
									__lowerCAmelCase											=				read_func(lowercase					,	lowercase	)

			with tempfile.TemporaryDirectory() as tmp_dir:
						# don't use fixed length for fair comparison
						# feats = datasets.Features(
						#     {"image": datasets.Sequence(datasets.Sequence(datasets.Value("float32"), SPEED_TEST_SHAPE[1]), SPEED_TEST_SHAPE[0])}
						# )
						__lowerCAmelCase											=				datasets.Features({"""image""": datasets.Sequence(datasets.Sequence(datasets.Value("""float32"""	)	)	)}	)
						__lowerCAmelCase											=				generate_examples(
						    features=lowercase					,	num_examples=lowercase					,	seq_shapes={"""image""": SPEED_TEST_SHAPE}	)
						__lowerCAmelCase											=				write(lowercase					,	lowercase					,	lowercase	)
						for read_func in read_functions:
									__lowerCAmelCase											=				read_func(lowercase					,	lowercase	)

			with tempfile.TemporaryDirectory() as tmp_dir:
						# don't use fixed length for fair comparison
						# feats = datasets.Features(
						#     {"image": datasets.Sequence(datasets.Value("float32"), SPEED_TEST_SHAPE[0] * SPEED_TEST_SHAPE[1])}
						# )
						__lowerCAmelCase											=				datasets.Features({"""image""": datasets.Sequence(datasets.Value("""float32"""	)	)}	)
						__lowerCAmelCase											=				generate_examples(
						    features=lowercase					,	num_examples=lowercase					,	seq_shapes={"""image""": [SPEED_TEST_SHAPE[0] * SPEED_TEST_SHAPE[1]]}					,	)
						__lowerCAmelCase											=				write(lowercase					,	lowercase					,	lowercase	)
						for read_func in read_functions:
									__lowerCAmelCase											=				read_func(lowercase					,	lowercase	)

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


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

