code: 
  
  
  
  
  """simple docstring"""
  
  import os
  import tempfile
  from pathlib import Path
  from unittest import TestCase
  
  import pyarrow as pa
  import pytest
  
  from datasets.arrow_dataset import Dataset
  from datasets.arrow_reader import ArrowReader, BaseReader, FileInstructions, ReadInstruction, make_file_instructions
  from datasets.info import DatasetInfo
  from datasets.splits import NamedSplit, Split, SplitDict, SplitInfo
  
  from .utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases
  class 			lowerCamelCase__    (  lowerCamelCase_			):
  
  
  			def 							lowerCamelCase_				(     self				, SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE=False							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					snake_case			,				snake_case			,				snake_case		:						List[Any]								=				(
  					    filename_skip_take["filename"],
  					    filename_skip_take["skip"] if "skip" in filename_skip_take else None,
  					    filename_skip_take["take"] if "take" in filename_skip_take else None,
  					)
  					open(os.path.join(SCREAMING_SNAKE_CASE							)				, "wb"							).close()
  					snake_case		:						Any								=				pa.Table.from_pydict({"filename": [Path(SCREAMING_SNAKE_CASE							).name] * 100}							)
  					if take == -1:
  							snake_case		:						Dict								=				len(SCREAMING_SNAKE_CASE							) - skip
  					if skip is not None and take is not None:
  							snake_case		:						Dict								=				pa_table.slice(SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE							)
  					return pa_table
  
  
  
  
  
  
  
  class 			lowerCamelCase__    (  lowerCamelCase_			):
  
  
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					snake_case		:						int								=				"my_name"
  					snake_case		:						Union[str, Any]								=				SplitInfo(name="train"				, num_examples=100							)
  					snake_case		:						Tuple								=				SplitInfo(name="test"				, num_examples=100							)
  					snake_case		:						int								=				[train_info, test_info]
  					snake_case		:						Dict								=				SplitDict()
  					split_dict.add(SCREAMING_SNAKE_CASE							)
  					split_dict.add(SCREAMING_SNAKE_CASE							)
  					snake_case		:						Union[str, Any]								=				DatasetInfo(splits=SCREAMING_SNAKE_CASE							)
  
  					with tempfile.TemporaryDirectory() as tmp_dir:
  							snake_case		:						List[Any]								=				ReaderTest(SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE							)
  
  							snake_case		:						List[str]								=				"test[:33%]"
  							snake_case		:						Tuple								=				Dataset(**reader.read(SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE							)							)
  							self.assertEqual(dset["filename"][0]				, F'''{name}-test'''							)
  							self.assertEqual(dset.num_rows				, 33							)
  							self.assertEqual(dset.num_columns				, 1							)
  
  							snake_case		:						Any								=				["train", "test[:33%]"]
  							snake_case		:						int								=				[Split.TRAIN, ReadInstruction.from_spec("test[:33%]"							)]
  							for instructions in [instructionsa, instructionsa]:
  									snake_case		:						Any								=				[reader.read(SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE							) for instr in instructions]
  									snake_case			,				snake_case		:						List[str]								=				(Dataset(**SCREAMING_SNAKE_CASE							) for dataset_kwargs in datasets_kwargs)
  									self.assertEqual(train_dset["filename"][0]				, F'''{name}-train'''							)
  									self.assertEqual(train_dset.num_rows				, 100							)
  									self.assertEqual(train_dset.num_columns				, 1							)
  									self.assertIsInstance(train_dset.split				, SCREAMING_SNAKE_CASE							)
  									self.assertEqual(str(train_dset.split							)				, "train"							)
  									self.assertEqual(test_dset["filename"][0]				, F'''{name}-test'''							)
  									self.assertEqual(test_dset.num_rows				, 33							)
  									self.assertEqual(test_dset.num_columns				, 1							)
  									self.assertIsInstance(test_dset.split				, SCREAMING_SNAKE_CASE							)
  									self.assertEqual(str(test_dset.split							)				, "test[:33%]"							)
  									del train_dset, test_dset
  
  
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					snake_case		:						Optional[int]								=				"my_name"
  					snake_case		:						List[str]								=				SplitInfo(name="train"				, num_examples=1_000				, shard_lengths=[100] * 10							)
  					snake_case		:						Union[str, Any]								=				[train_info]
  					snake_case		:						List[str]								=				SplitDict()
  					split_dict.add(SCREAMING_SNAKE_CASE							)
  					snake_case		:						Optional[int]								=				DatasetInfo(splits=SCREAMING_SNAKE_CASE							)
  
  					with tempfile.TemporaryDirectory() as tmp_dir:
  							snake_case		:						Tuple								=				ReaderTest(SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE							)
  
  							snake_case		:						int								=				"train[:33%]"
  							snake_case		:						int								=				Dataset(**reader.read(SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE							)							)
  							self.assertEqual(dset["filename"][0]				, F'''{name}-train-00000-of-00010'''							)
  							self.assertEqual(dset["filename"][-1]				, F'''{name}-train-00003-of-00010'''							)
  							self.assertEqual(dset.num_rows				, 330							)
  							self.assertEqual(dset.num_columns				, 1							)
  
  
  
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					snake_case		:						Optional[Any]								=				SplitInfo(name="train"				, num_examples=100							)
  					snake_case		:						Dict								=				SplitInfo(name="test"				, num_examples=100							)
  					snake_case		:						str								=				SplitDict()
  					split_dict.add(SCREAMING_SNAKE_CASE							)
  					split_dict.add(SCREAMING_SNAKE_CASE							)
  					snake_case		:						Dict								=				DatasetInfo(splits=SCREAMING_SNAKE_CASE							)
  
  					with tempfile.TemporaryDirectory() as tmp_dir:
  							snake_case		:						str								=				ReaderTest(SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE							)
  
  							snake_case		:						int								=				[
  							    {"filename": os.path.join(SCREAMING_SNAKE_CASE				, "train"							)},
  							    {"filename": os.path.join(SCREAMING_SNAKE_CASE				, "test"							), "skip": 10, "take": 10},
  							]
  							snake_case		:						Union[str, Any]								=				Dataset(**reader.read_files(SCREAMING_SNAKE_CASE				, original_instructions="train+test[10:20]"							)							)
  							self.assertEqual(dset.num_rows				, 110							)
  							self.assertEqual(dset.num_columns				, 1							)
  							del dset
  
  
  
  
  
  
  
  @pytest.mark.parametrize("in_memory"       ,     [False, True]   )
  def   UpperCamelCase__     (		lowercase__  :     Optional[Any]       ,     lowercase__  :     List[Any]       ,     lowercase__  :     Optional[int]   ):
  		snake_case		:						Dict								=				arrow_file
  		with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
  				snake_case		:						Any								=				ArrowReader.read_table(lowercase__       ,     in_memory=lowercase__   )
  		assert table.shape == dataset.data.shape
  		assert set(table.column_names   ) == set(dataset.data.column_names   )
  		assert dict(table.to_pydict()   ) == dict(dataset.data.to_pydict()   )  # to_pydict returns OrderedDict
  
  
  
  @pytest.mark.parametrize("in_memory"       ,     [False, True]   )
  def   UpperCamelCase__     (		lowercase__  :     List[str]       ,     lowercase__  :     str       ,     lowercase__  :     str   ):
  		snake_case		:						str								=				arrow_file
  		snake_case		:						Any								=				ArrowReader(""       ,     lowercase__   )
  		with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
  				snake_case		:						List[Any]								=				reader.read_files([{"filename": filename}]       ,     in_memory=lowercase__   )
  		assert dataset_kwargs.keys() == {"arrow_table", "info", "split"}
  		snake_case		:						List[str]								=				dataset_kwargs["arrow_table"]
  		assert table.shape == dataset.data.shape
  		assert set(table.column_names   ) == set(dataset.data.column_names   )
  		assert dict(table.to_pydict()   ) == dict(dataset.data.to_pydict()   )  # to_pydict returns OrderedDict
  
  
  
  def   UpperCamelCase__     (		):
  		assert ReadInstruction("train"       ,     to=10       ,     unit="abs"   ).to_spec() == "train[:10]"
  		assert ReadInstruction("train"       ,     from_=-80       ,     to=10       ,     unit="%"   ).to_spec() == "train[-80%:10%]"
  
  		snake_case		:						Tuple								=				"train+test"
  		assert ReadInstruction.from_spec(lowercase__   ).to_spec() == spec_train_test
  
  		snake_case		:						str								=				"train[2:10]"
  		assert ReadInstruction.from_spec(lowercase__   ).to_spec() == spec_train_abs
  
  		snake_case		:						Optional[Any]								=				"train[15%:-20%]"
  		assert ReadInstruction.from_spec(lowercase__   ).to_spec() == spec_train_pct
  
  		snake_case		:						Union[str, Any]								=				"train[:10%](closest)"
  		assert ReadInstruction.from_spec(lowercase__   ).to_spec() == "train[:10%]"
  
  		snake_case		:						Dict								=				"train[:10%](pct1_dropremainder)"
  		assert ReadInstruction.from_spec(lowercase__   ).to_spec() == spec_train_pct_rounding
  
  		snake_case		:						Tuple								=				"train[:10%](pct1_dropremainder)+test[-10%:](pct1_dropremainder)"
  		assert ReadInstruction.from_spec(lowercase__   ).to_spec() == spec_train_test_pct_rounding
  
  
  
  def   UpperCamelCase__     (		):
  		snake_case		:						List[str]								=				"dummy"
  		snake_case		:						Optional[Any]								=				[SplitInfo(name="train"       ,     num_examples=100   )]
  		snake_case		:						List[str]								=				"train[:33%]"
  		snake_case		:						Optional[int]								=				"arrow"
  		snake_case		:						Union[str, Any]								=				"prefix"
  
  		snake_case		:						Tuple								=				make_file_instructions(lowercase__       ,     lowercase__       ,     lowercase__       ,     lowercase__       ,     lowercase__   )
  		assert isinstance(lowercase__       ,     lowercase__   )
  		assert file_instructions.num_examples == 33
  		assert file_instructions.file_instructions == [
  		    {"filename": os.path.join(lowercase__       ,     F'''{name}-train.arrow'''   ), "skip": 0, "take": 33}
  		]
  
  		snake_case		:						Tuple								=				[SplitInfo(name="train"       ,     num_examples=100       ,     shard_lengths=[10] * 10   )]
  		snake_case		:						List[str]								=				make_file_instructions(lowercase__       ,     lowercase__       ,     lowercase__       ,     lowercase__       ,     lowercase__   )
  		assert isinstance(lowercase__       ,     lowercase__   )
  		assert file_instructions.num_examples == 33
  		assert file_instructions.file_instructions == [
  		    {"filename": os.path.join(lowercase__       ,     F'''{name}-train-00000-of-00010.arrow'''   ), "skip": 0, "take": -1},
  		    {"filename": os.path.join(lowercase__       ,     F'''{name}-train-00001-of-00010.arrow'''   ), "skip": 0, "take": -1},
  		    {"filename": os.path.join(lowercase__       ,     F'''{name}-train-00002-of-00010.arrow'''   ), "skip": 0, "take": -1},
  		    {"filename": os.path.join(lowercase__       ,     F'''{name}-train-00003-of-00010.arrow'''   ), "skip": 0, "take": 3},
  		]
  
  
  
  @pytest.mark.parametrize("name, expected_exception"       ,     [(None, TypeError), ("", ValueError)]   )
  def   UpperCamelCase__     (		lowercase__  :     Any       ,     lowercase__  :     List[Any]   ):
  		snake_case		:						Dict								=				[SplitInfo(name="train"       ,     num_examples=100   )]
  		snake_case		:						Any								=				"train"
  		snake_case		:						List[str]								=				"arrow"
  		snake_case		:						Any								=				"prefix_path"
  		with pytest.raises(lowercase__   ):
  				snake_case		:						Optional[int]								=				make_file_instructions(lowercase__       ,     lowercase__       ,     lowercase__       ,     lowercase__       ,     lowercase__   )
  
  
style_context: 
  
  
  
  
  """simple docstring"""
  
  import gc
  import unittest
  
  import numpy as np
  import torch
  
  from diffusers import RePaintPipeline, RePaintScheduler, UNetaDModel
  from diffusers.utils.testing_utils import (
      enable_full_determinism,
      load_image,
      load_numpy,
      nightly,
      require_torch_gpu,
      skip_mps,
      torch_device,
  )
  
  from ..pipeline_params import IMAGE_INPAINTING_BATCH_PARAMS, IMAGE_INPAINTING_PARAMS
  from ..test_pipelines_common import PipelineTesterMixin
  
  
  enable_full_determinism()
  class 			lowerCamelCase__    (  lowerCamelCase_				,  unittest.TestCase			):
  			a__ :      Optional[Any]								=     RePaintPipeline
  			a__ :      List[Any]								=     IMAGE_INPAINTING_PARAMS - {"""width""", """height""", """guidance_scale"""}
  			a__ :      List[Any]								=     PipelineTesterMixin.required_optional_params - {
  			    """latents""",
  			    """num_images_per_prompt""",
  			    """callback""",
  			    """callback_steps""",
  			}
  			a__ :      List[Any]								=     IMAGE_INPAINTING_BATCH_PARAMS
  			a__ :      List[str]								=     False
  
  
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					torch.manual_seed(0							)
  					torch.manual_seed(0							)
  					snake_case		:						Union[str, Any]								=				UNetaDModel(
  					    block_out_channels=(32, 64)				, layers_per_block=2				, sample_size=32				, in_channels=3				, out_channels=3				, down_block_types=("DownBlock2D", "AttnDownBlock2D")				, up_block_types=("AttnUpBlock2D", "UpBlock2D")				, )
  					snake_case		:						Union[str, Any]								=				RePaintScheduler()
  					snake_case		:						Optional[Any]								=				{"unet": unet, "scheduler": scheduler}
  					return components
  
  
  			def 							lowerCamelCase_				(     self				, SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE=0							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					if str(SCREAMING_SNAKE_CASE							).startswith("mps"							):
  							snake_case		:						str								=				torch.manual_seed(SCREAMING_SNAKE_CASE							)
  					else:
  							snake_case		:						List[Any]								=				torch.Generator(device=SCREAMING_SNAKE_CASE							).manual_seed(SCREAMING_SNAKE_CASE							)
  					snake_case		:						Any								=				np.random.RandomState(SCREAMING_SNAKE_CASE							).standard_normal((1, 3, 32, 32)							)
  					snake_case		:						Any								=				torch.from_numpy(SCREAMING_SNAKE_CASE							).to(device=SCREAMING_SNAKE_CASE				, dtype=torch.floataa							)
  					snake_case		:						List[str]								=				(image > 0).to(device=SCREAMING_SNAKE_CASE				, dtype=torch.floataa							)
  					snake_case		:						int								=				{
  					    "image": image,
  					    "mask_image": mask,
  					    "generator": generator,
  					    "num_inference_steps": 5,
  					    "eta": 0.0,
  					    "jump_length": 2,
  					    "jump_n_sample": 2,
  					    "output_type": "numpy",
  					}
  					return inputs
  
  
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					snake_case		:						List[Any]								=				"cpu"  # ensure determinism for the device-dependent torch.Generator
  					snake_case		:						Any								=				self.get_dummy_components()
  					snake_case		:						Optional[int]								=				RePaintPipeline(**SCREAMING_SNAKE_CASE							)
  					snake_case		:						Optional[Any]								=				sd_pipe.to(SCREAMING_SNAKE_CASE							)
  					sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE							)
  
  					snake_case		:						Any								=				self.get_dummy_inputs(SCREAMING_SNAKE_CASE							)
  					snake_case		:						Dict								=				sd_pipe(**SCREAMING_SNAKE_CASE							).images
  					snake_case		:						Optional[int]								=				image[0, -3:, -3:, -1]
  
  					assert image.shape == (1, 32, 32, 3)
  					snake_case		:						List[str]								=				np.array([1.00_00, 0.54_26, 0.54_97, 0.22_00, 1.00_00, 1.00_00, 0.56_23, 1.00_00, 0.62_74]							)
  
  					assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-3
  
  
  			@skip_mps
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					return super().test_save_load_local()
  
  
  			@unittest.skip("non-deterministic pipeline"							)
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					return super().test_inference_batch_single_identical()
  
  
  			@skip_mps
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					return super().test_dict_tuple_outputs_equivalent()
  
  
  			@skip_mps
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					return super().test_save_load_optional_components()
  
  
  
  			@skip_mps
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					return super().test_attention_slicing_forward_pass()
  
  
  
  
  
  
  
  @nightly
  @require_torch_gpu
  class 			lowerCamelCase__    (  unittest.TestCase			):
  
  
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					super().tearDown()
  					gc.collect()
  					torch.cuda.empty_cache()
  
  
  
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					snake_case		:						Optional[int]								=				load_image(
  					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/"
  					    "repaint/celeba_hq_256.png"							)
  					snake_case		:						Any								=				load_image(
  					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/mask_256.png"							)
  					snake_case		:						List[Any]								=				load_numpy(
  					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/"
  					    "repaint/celeba_hq_256_result.npy"							)
  
  					snake_case		:						str								=				"google/ddpm-ema-celebahq-256"
  					snake_case		:						Optional[Any]								=				UNetaDModel.from_pretrained(SCREAMING_SNAKE_CASE							)
  					snake_case		:						Optional[int]								=				RePaintScheduler.from_pretrained(SCREAMING_SNAKE_CASE							)
  
  					snake_case		:						str								=				RePaintPipeline(unet=SCREAMING_SNAKE_CASE				, scheduler=SCREAMING_SNAKE_CASE							).to(SCREAMING_SNAKE_CASE							)
  					repaint.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE							)
  					repaint.enable_attention_slicing()
  
  					snake_case		:						Union[str, Any]								=				torch.manual_seed(0							)
  					snake_case		:						int								=				repaint(
  					    SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE				, num_inference_steps=250				, eta=0.0				, jump_length=10				, jump_n_sample=10				, generator=SCREAMING_SNAKE_CASE				, output_type="np"				, )
  					snake_case		:						List[Any]								=				output.images[0]
  
  					assert image.shape == (256, 256, 3)
  					assert np.abs(expected_image - image							).mean() < 1E-2
  
  