






# Lint as: python3
# pylint: enable=line-too-long
# pylint: disable=g-import-not-at-top,g-bad-import-order,wrong-import-position

UpperCAmelCase__				   =		"2.13.1"

import platform

import pyarrow
from packaging import version


if version.parse(platform.python_version()) < version.parse("3.7"):
     raise ImportWarning(
         "To use `datasets`, Python>=3.7 is required, and the current version of Python doesn't match this condition."
     )

if version.parse(pyarrow.__version__).major < 8:
     raise ImportWarning(
         "To use `datasets`, the module `pyarrow>=8.0.0` is required, and the current version of `pyarrow` doesn't match this condition.\n"
         "If you are running this in a Google Colab, you should probably just restart the runtime to use the right version of `pyarrow`."
     )

del platform
del pyarrow
del version

from .arrow_dataset import Dataset
from .arrow_reader import ReadInstruction
from .builder import ArrowBasedBuilder, BeamBasedBuilder, BuilderConfig, DatasetBuilder, GeneratorBasedBuilder
from .combine import concatenate_datasets, interleave_datasets
from .dataset_dict import DatasetDict, IterableDatasetDict
from .download import *
from .features import *
from .fingerprint import disable_caching, enable_caching, is_caching_enabled, set_caching_enabled
from .info import DatasetInfo, MetricInfo
from .inspect import (
    get_dataset_config_info,
    get_dataset_config_names,
    get_dataset_infos,
    get_dataset_split_names,
    inspect_dataset,
    inspect_metric,
    list_datasets,
    list_metrics,
)
from .iterable_dataset import IterableDataset
from .load import load_dataset, load_dataset_builder, load_from_disk, load_metric
from .metric import Metric
from .splits import (
    NamedSplit,
    NamedSplitAll,
    Split,
    SplitBase,
    SplitDict,
    SplitGenerator,
    SplitInfo,
    SubSplitInfo,
    percent,
)
from .tasks import *
from .utils import *
from .utils import logging


# deprecated modules
from datasets import arrow_dataset as _arrow_dataset  # isort:skip
from datasets import utils as _utils  # isort:skip
from datasets.utils import download_manager as _deprecated_download_manager  # isort:skip

UpperCAmelCase__				   =		concatenate_datasets
UpperCAmelCase__				   =		DownloadConfig
UpperCAmelCase__				   =		DownloadManager
UpperCAmelCase__				   =		DownloadMode
UpperCAmelCase__				   =		DownloadConfig
UpperCAmelCase__				   =		DownloadMode
UpperCAmelCase__				   =		DownloadManager

del _arrow_dataset, _utils, _deprecated_download_manager








import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, PNDMScheduler, StableDiffusionInpaintPipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class     lowerCamelCase__(					__lowerCamelCase ,     __lowerCamelCase ,     __lowerCamelCase ,     unittest.TestCase):
      UpperCAmelCase__      :	Union[str, Any]    		=   StableDiffusionInpaintPipeline
      UpperCAmelCase__      :	Dict    		=   TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
      UpperCAmelCase__      :	str    		=   TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
      UpperCAmelCase__      :	int    		=   frozenset(
          [])  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
      UpperCAmelCase__      :	Union[str, Any]    		=   frozenset([])
      def 			lowerCAmelCase__						(				self: str   ):
             torch.manual_seed(0   )
             __lowerCamelCase   =						UNetaDConditionModel(
                 block_out_channels=(32, 64)    ,       layers_per_block=2    ,       sample_size=32    ,       in_channels=9    ,       out_channels=4    ,       down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")    ,       up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")    ,       cross_attention_dim=32    ,       attention_head_dim=(2, 4)    ,       use_linear_projection=UpperCamelCase_    ,       )
             __lowerCamelCase   =						PNDMScheduler(skip_prk_steps=UpperCamelCase_   )
             torch.manual_seed(0   )
             __lowerCamelCase   =						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    ,       sample_size=1_28    ,       )
             torch.manual_seed(0   )
             __lowerCamelCase   =						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=10_00    ,       hidden_act="""gelu"""    ,       projection_dim=5_12    ,       )
             __lowerCamelCase   =						CLIPTextModel(UpperCamelCase_   )
             __lowerCamelCase   =						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""   )

             __lowerCamelCase   =						{
                 """unet""": unet,
                 """scheduler""": scheduler,
                 """vae""": vae,
                 """text_encoder""": text_encoder,
                 """tokenizer""": tokenizer,
                 """safety_checker""": None,
                 """feature_extractor""": None,
             }
             return components
      def 			lowerCAmelCase__						(				self: List[Any]    ,       UpperCamelCase_: Any    ,       UpperCamelCase_: List[Any]=0   ):
             # TODO: use tensor inputs instead of PIL, this is here just to leave the old expected_slices untouched
             __lowerCamelCase   =						floats_tensor((1, 3, 32, 32)    ,       rng=random.Random(UpperCamelCase_   )   ).to(UpperCamelCase_   )
             __lowerCamelCase   =						image.cpu().permute(0    ,       2    ,       3    ,       1   )[0]
             __lowerCamelCase   =						Image.fromarray(np.uinta(UpperCamelCase_   )   ).convert("""RGB"""   ).resize((64, 64)   )
             __lowerCamelCase   =						Image.fromarray(np.uinta(image + 4   )   ).convert("""RGB"""   ).resize((64, 64)   )
             if str(UpperCamelCase_   ).startswith("""mps"""   ):
                    __lowerCamelCase   =						torch.manual_seed(UpperCamelCase_   )
             else:
                    __lowerCamelCase   =						torch.Generator(device=UpperCamelCase_   ).manual_seed(UpperCamelCase_   )
             __lowerCamelCase   =						{
                 """prompt""": """A painting of a squirrel eating a burger""",
                 """image""": init_image,
                 """mask_image""": mask_image,
                 """generator""": generator,
                 """num_inference_steps""": 2,
                 """guidance_scale""": 6.0,
                 """output_type""": """numpy""",
             }
             return inputs
      def 			lowerCAmelCase__						(				self: str   ):
             __lowerCamelCase   =						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
             __lowerCamelCase   =						self.get_dummy_components()
             __lowerCamelCase   =						StableDiffusionInpaintPipeline(**UpperCamelCase_   )
             __lowerCamelCase   =						sd_pipe.to(UpperCamelCase_   )
             sd_pipe.set_progress_bar_config(disable=UpperCamelCase_   )

             __lowerCamelCase   =						self.get_dummy_inputs(UpperCamelCase_   )
             __lowerCamelCase   =						sd_pipe(**UpperCamelCase_   ).images
             __lowerCamelCase   =						image[0, -3:, -3:, -1]

             assert image.shape == (1, 64, 64, 3)
             __lowerCamelCase   =						np.array([0.4727, 0.5735, 0.3941, 0.5446, 0.5926, 0.4394, 0.5062, 0.4654, 0.4476]   )

             assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1E-2




      def 			lowerCAmelCase__						(				self: int   ):
             super().test_inference_batch_single_identical(expected_max_diff=3E-3   )





@slow
@require_torch_gpu
class     lowerCamelCase__(					unittest.TestCase):
      def 			lowerCAmelCase__						(				self: str   ):
             # clean up the VRAM after each test
             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()
      def 			lowerCAmelCase__						(				self: List[Any]   ):
             __lowerCamelCase   =						load_image(
                 """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                 """/sd2-inpaint/init_image.png"""   )
             __lowerCamelCase   =						load_image(
                 """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"""   )
             __lowerCamelCase   =						load_numpy(
                 """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"""
                 """/yellow_cat_sitting_on_a_park_bench.npy"""   )

             __lowerCamelCase   =						"""stabilityai/stable-diffusion-2-inpainting"""
             __lowerCamelCase   =						StableDiffusionInpaintPipeline.from_pretrained(UpperCamelCase_    ,       safety_checker=UpperCamelCase_   )
             pipe.to(UpperCamelCase_   )
             pipe.set_progress_bar_config(disable=UpperCamelCase_   )
             pipe.enable_attention_slicing()

             __lowerCamelCase   =						"""Face of a yellow cat, high resolution, sitting on a park bench"""

             __lowerCamelCase   =						torch.manual_seed(0   )
             __lowerCamelCase   =						pipe(
                 prompt=UpperCamelCase_    ,       image=UpperCamelCase_    ,       mask_image=UpperCamelCase_    ,       generator=UpperCamelCase_    ,       output_type="""np"""    ,       )
             __lowerCamelCase   =						output.images[0]

             assert image.shape == (5_12, 5_12, 3)
             assert np.abs(expected_image - image   ).max() < 9E-3
      def 			lowerCAmelCase__						(				self: Optional[int]   ):
             __lowerCamelCase   =						load_image(
                 """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                 """/sd2-inpaint/init_image.png"""   )
             __lowerCamelCase   =						load_image(
                 """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"""   )
             __lowerCamelCase   =						load_numpy(
                 """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"""
                 """/yellow_cat_sitting_on_a_park_bench_fp16.npy"""   )

             __lowerCamelCase   =						"""stabilityai/stable-diffusion-2-inpainting"""
             __lowerCamelCase   =						StableDiffusionInpaintPipeline.from_pretrained(
                 UpperCamelCase_    ,       torch_dtype=torch.floataa    ,       safety_checker=UpperCamelCase_    ,       )
             pipe.to(UpperCamelCase_   )
             pipe.set_progress_bar_config(disable=UpperCamelCase_   )
             pipe.enable_attention_slicing()

             __lowerCamelCase   =						"""Face of a yellow cat, high resolution, sitting on a park bench"""

             __lowerCamelCase   =						torch.manual_seed(0   )
             __lowerCamelCase   =						pipe(
                 prompt=UpperCamelCase_    ,       image=UpperCamelCase_    ,       mask_image=UpperCamelCase_    ,       generator=UpperCamelCase_    ,       output_type="""np"""    ,       )
             __lowerCamelCase   =						output.images[0]

             assert image.shape == (5_12, 5_12, 3)
             assert np.abs(expected_image - image   ).max() < 5E-1




      def 			lowerCAmelCase__						(				self: int   ):
             torch.cuda.empty_cache()
             torch.cuda.reset_max_memory_allocated()
             torch.cuda.reset_peak_memory_stats()

             __lowerCamelCase   =						load_image(
                 """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                 """/sd2-inpaint/init_image.png"""   )
             __lowerCamelCase   =						load_image(
                 """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"""   )

             __lowerCamelCase   =						"""stabilityai/stable-diffusion-2-inpainting"""
             __lowerCamelCase   =						PNDMScheduler.from_pretrained(UpperCamelCase_    ,       subfolder="""scheduler"""   )
             __lowerCamelCase   =						StableDiffusionInpaintPipeline.from_pretrained(
                 UpperCamelCase_    ,       safety_checker=UpperCamelCase_    ,       scheduler=UpperCamelCase_    ,       torch_dtype=torch.floataa    ,       )
             pipe.to(UpperCamelCase_   )
             pipe.set_progress_bar_config(disable=UpperCamelCase_   )
             pipe.enable_attention_slicing(1   )
             pipe.enable_sequential_cpu_offload()

             __lowerCamelCase   =						"""Face of a yellow cat, high resolution, sitting on a park bench"""

             __lowerCamelCase   =						torch.manual_seed(0   )
             __lowerCamelCase   =						pipe(
                 prompt=UpperCamelCase_    ,       image=UpperCamelCase_    ,       mask_image=UpperCamelCase_    ,       generator=UpperCamelCase_    ,       num_inference_steps=2    ,       output_type="""np"""    ,       )

             __lowerCamelCase   =						torch.cuda.max_memory_allocated()
             # make sure that less than 2.65 GB is allocated
             assert mem_bytes < 2.65 * 10**9

