

import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


__UpperCAmelCase    				=     datasets.utils.logging.get_logger(__name__)





@dataclass
class A__    ( datasets.BuilderConfig ):



 """simple docstring"""







 _lowercase     :					Optional[datasets.Features]      =						None
 _lowercase     :					str      =						"utf-8"
 _lowercase     :					Optional[str]      =						None
 _lowercase     :					Optional[str]      =						None
 _lowercase     :					bool      =						True  # deprecated
 _lowercase     :					Optional[int]      =						None  # deprecated
 _lowercase     :					int      =						10 << 20  # 10MB
 _lowercase     :					Optional[bool]      =						None





class A__    ( datasets.ArrowBasedBuilder ):



 """simple docstring"""







 _lowercase     :					Tuple      =						JsonConfig

 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      if self.config.block_size is not None:
           logger.warning("The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead"       )
           _lowerCAmelCase   :					Tuple        =    self.config.block_size
      if self.config.use_threads is not True:
           logger.warning(
               "The JSON loader parameter `use_threads` is deprecated and doesn't have any effect anymore."       )
      if self.config.newlines_in_values is not None:
           raise ValueError("The JSON loader parameter `newlines_in_values` is no longer supported"       )
      return datasets.DatasetInfo(features=self.config.features       )

 def    __magic_name__       (		self :	Dict     ,						A_ :	int       ):






      '''simple docstring'''




      if not self.config.data_files:
           raise ValueError(F'''At least one data file must be specified, but got data_files={self.config.data_files}'''       )
      _lowerCAmelCase   :					Tuple        =    dl_manager.download_and_extract(self.config.data_files       )
      if isinstance(A_     ,						(str, list, tuple)       ):
           _lowerCAmelCase   :					int        =    data_files
           if isinstance(A_     ,						A_       ):
                _lowerCAmelCase   :					Optional[int]        =    [files]
           _lowerCAmelCase   :					List[Any]        =    [dl_manager.iter_files(A_       ) for file in files]
           return [datasets.SplitGenerator(name=datasets.Split.TRAIN     ,						gen_kwargs={"files": files}       )]
      _lowerCAmelCase   :					Any        =    []
      for split_name, files in data_files.items():
           if isinstance(A_     ,						A_       ):
                _lowerCAmelCase   :					int        =    [files]
           _lowerCAmelCase   :					Tuple        =    [dl_manager.iter_files(A_       ) for file in files]
           splits.append(datasets.SplitGenerator(name=A_     ,						gen_kwargs={"files": files}       )       )
      return splits

 def    __magic_name__       (		self :	List[str]     ,						A_ :	pa.Table       ):






      '''simple docstring'''




      if self.config.features is not None:
           # adding missing columns
           for column_name in set(self.config.features       ) - set(pa_table.column_names       ):
                _lowerCAmelCase   :					Union[str, Any]        =    self.config.features.arrow_schema.field(A_       ).type
                _lowerCAmelCase   :					int        =    pa_table.append_column(A_     ,						pa.array([None] * len(A_       )     ,						type=A_       )       )
           # more expensive cast to support nested structures with keys in a different order
           # allows str <-> int/float or str to Audio for example
           _lowerCAmelCase   :					Tuple        =    table_cast(A_     ,						self.config.features.arrow_schema       )
      return pa_table


 def    __magic_name__       (		self :	Tuple     ,						A_ :	Dict       ):






      '''simple docstring'''




      for file_idx, file in enumerate(itertools.chain.from_iterable(A_       )       ):
           # If the file is one json object and if we need to look at the list of items in one specific field
           if self.config.field is not None:
                with open(A_     ,						encoding=self.config.encoding     ,						errors=self.config.encoding_errors       ) as f:
                     _lowerCAmelCase   :					Optional[Any]        =    json.load(A_       )

                # We keep only the field we are interested in
                _lowerCAmelCase   :					Tuple        =    dataset[self.config.field]

                # We accept two format: a list of dicts or a dict of lists
                if isinstance(A_     ,						(list, tuple)       ):
                     _lowerCAmelCase   :					Tuple        =    set().union(*[row.keys() for row in dataset]       )
                     _lowerCAmelCase   :					Union[str, Any]        =    {col: [row.get(A_       ) for row in dataset] for col in keys}
                else:
                     _lowerCAmelCase   :					Tuple        =    dataset
                _lowerCAmelCase   :					List[str]        =    pa.Table.from_pydict(A_       )
                yield file_idx, self._cast_table(A_       )

           # If the file has one json object per line
           else:
                with open(A_     ,						"rb"       ) as f:
                     _lowerCAmelCase   :					Union[str, Any]        =    0
                     # Use block_size equal to the chunk size divided by 32 to leverage multithreading
                     # Set a default minimum value of 16kB if the chunk size is really small
                     _lowerCAmelCase   :					int        =    max(self.config.chunksize // 3_2     ,						1_6 << 1_0       )
                     _lowerCAmelCase   :					List[Any]        =    (
                         self.config.encoding_errors if self.config.encoding_errors is not None else "strict"
                     )
                     while True:
                          _lowerCAmelCase   :					str        =    f.read(self.config.chunksize       )
                          if not batch:
                               break
                          # Finish current line
                          try:
                               batch += f.readline()
                          except (AttributeError, io.UnsupportedOperation):
                               batch += readline(A_       )
                          # PyArrow only accepts utf-8 encoded bytes
                          if self.config.encoding != "utf-8":
                               _lowerCAmelCase   :					int        =    batch.decode(self.config.encoding     ,						errors=A_       ).encode("utf-8"       )
                          try:
                               while True:
                                    try:
                                         _lowerCAmelCase   :					Tuple        =    paj.read_json(
                                             io.BytesIO(A_       )     ,						read_options=paj.ReadOptions(block_size=A_       )       )
                                         break
                                    except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
                                         if (
                                             isinstance(A_     ,						pa.ArrowInvalid       )
                                             and "straddling" not in str(A_       )
                                             or block_size > len(A_       )
                                         ):
                                              raise
                                         else:
                                              # Increase the block size in case it was too small.
                                              # The block size will be reset for the next file.
                                              logger.debug(
                                                  F'''Batch of {len(A_       )} bytes couldn\'t be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}.'''       )
                                              block_size *= 2
                          except pa.ArrowInvalid as e:
                               try:
                                    with open(
                                        A_     ,						encoding=self.config.encoding     ,						errors=self.config.encoding_errors       ) as f:
                                         _lowerCAmelCase   :					Tuple        =    json.load(A_       )
                               except json.JSONDecodeError:
                                    logger.error(F'''Failed to read file \'{file}\' with error {type(A_       )}: {e}'''       )
                                    raise e
                               # If possible, parse the file as a list of json objects and exit the loop
                               if isinstance(A_     ,						A_       ):  # list is the only sequence type supported in JSON
                                    try:
                                         _lowerCAmelCase   :					Dict        =    set().union(*[row.keys() for row in dataset]       )
                                         _lowerCAmelCase   :					List[str]        =    {col: [row.get(A_       ) for row in dataset] for col in keys}
                                         _lowerCAmelCase   :					Dict        =    pa.Table.from_pydict(A_       )
                                    except (pa.ArrowInvalid, AttributeError) as e:
                                         logger.error(F'''Failed to read file \'{file}\' with error {type(A_       )}: {e}'''       )
                                         raise ValueError(F'''Not able to read records in the JSON file at {file}.'''       ) from None
                                    yield file_idx, self._cast_table(A_       )
                                    break
                               else:
                                    logger.error(F'''Failed to read file \'{file}\' with error {type(A_       )}: {e}'''       )
                                    raise ValueError(
                                        F'''Not able to read records in the JSON file at {file}. '''
                                        F'''You should probably indicate the field of the JSON file containing your records. '''
                                        F'''This JSON file contain the following fields: {str(list(dataset.keys()       )       )}. '''
                                        F'''Select the correct one and provide it as `field=\'XXX\'` to the dataset loading method. '''       ) from None
                        # Uncomment for debugging (will print the Arrow table size and elements)
                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                          yield (file_idx, batch_idx), self._cast_table(A_       )
                          batch_idx += 1









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, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()





class A__    ( unittest.TestCase ):



 """simple docstring"""







 def    __magic_name__       (		self :	Any       ):






      '''simple docstring'''




      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()

 @property
 def    __magic_name__       (		self :	List[Any]       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Optional[int]        =    1
      _lowerCAmelCase   :					List[str]        =    3
      _lowerCAmelCase   :					int        =    (3_2, 3_2)

      _lowerCAmelCase   :					Dict        =    floats_tensor((batch_size, num_channels) + sizes     ,						rng=random.Random(0       )       ).to(A_       )
      return image

 @property
 def    __magic_name__       (		self :	List[str]       ):






      '''simple docstring'''




      torch.manual_seed(0       )
      _lowerCAmelCase   :					Dict        =    UNetaDConditionModel(
          block_out_channels=(3_2, 3_2, 6_4)     ,						layers_per_block=2     ,						sample_size=3_2     ,						in_channels=7     ,						out_channels=4     ,						down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D")     ,						up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D")     ,						cross_attention_dim=3_2     ,						attention_head_dim=8     ,						use_linear_projection=A_     ,						only_cross_attention=(True, True, False)     ,						num_class_embeds=1_0_0     ,						)
      return model

 @property
 def    __magic_name__       (		self :	Union[str, Any]       ):






      '''simple docstring'''




      torch.manual_seed(0       )
      _lowerCAmelCase   :					Union[str, Any]        =    AutoencoderKL(
          block_out_channels=[3_2, 3_2, 6_4]     ,						in_channels=3     ,						out_channels=3     ,						down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"]     ,						up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"]     ,						latent_channels=4     ,						)
      return model

 @property
 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      torch.manual_seed(0       )
      _lowerCAmelCase   :					Optional[Any]        =    CLIPTextConfig(
          bos_token_id=0     ,						eos_token_id=2     ,						hidden_size=3_2     ,						intermediate_size=3_7     ,						layer_norm_eps=1E-05     ,						num_attention_heads=4     ,						num_hidden_layers=5     ,						pad_token_id=1     ,						vocab_size=1_0_0_0     ,						hidden_act="gelu"     ,						projection_dim=5_1_2     ,						)
      return CLIPTextModel(A_       )

 def    __magic_name__       (		self :	str       ):






      '''simple docstring'''




      _lowerCAmelCase   :					List[str]        =    "cpu"  # ensure determinism for the device-dependent torch.Generator
      _lowerCAmelCase   :					Tuple        =    self.dummy_cond_unet_upscale
      _lowerCAmelCase   :					Optional[int]        =    DDPMScheduler()
      _lowerCAmelCase   :					Optional[Any]        =    DDIMScheduler(prediction_type="v_prediction"       )
      _lowerCAmelCase   :					str        =    self.dummy_vae
      _lowerCAmelCase   :					Any        =    self.dummy_text_encoder
      _lowerCAmelCase   :					Dict        =    CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

      _lowerCAmelCase   :					List[Any]        =    self.dummy_image.cpu().permute(0     ,						2     ,						3     ,						1       )[0]
      _lowerCAmelCase   :					List[Any]        =    Image.fromarray(np.uinta(A_       )       ).convert("RGB"       ).resize((6_4, 6_4)       )

      # make sure here that pndm scheduler skips prk
      _lowerCAmelCase   :					List[Any]        =    StableDiffusionUpscalePipeline(
          unet=A_     ,						low_res_scheduler=A_     ,						scheduler=A_     ,						vae=A_     ,						text_encoder=A_     ,						tokenizer=A_     ,						max_noise_level=3_5_0     ,						)
      _lowerCAmelCase   :					Optional[int]        =    sd_pipe.to(A_       )
      sd_pipe.set_progress_bar_config(disable=A_       )

      _lowerCAmelCase   :					List[str]        =    "A painting of a squirrel eating a burger"
      _lowerCAmelCase   :					Tuple        =    torch.Generator(device=A_       ).manual_seed(0       )
      _lowerCAmelCase   :					Optional[int]        =    sd_pipe(
          [prompt]     ,						image=A_     ,						generator=A_     ,						guidance_scale=6.0     ,						noise_level=2_0     ,						num_inference_steps=2     ,						output_type="np"     ,						)

      _lowerCAmelCase   :					Tuple        =    output.images

      _lowerCAmelCase   :					Optional[int]        =    torch.Generator(device=A_       ).manual_seed(0       )
      _lowerCAmelCase   :					Union[str, Any]        =    sd_pipe(
          [prompt]     ,						image=A_     ,						generator=A_     ,						guidance_scale=6.0     ,						noise_level=2_0     ,						num_inference_steps=2     ,						output_type="np"     ,						return_dict=A_     ,						)[0]

      _lowerCAmelCase   :					List[Any]        =    image[0, -3:, -3:, -1]
      _lowerCAmelCase   :					List[str]        =    image_from_tuple[0, -3:, -3:, -1]

      _lowerCAmelCase   :					Optional[int]        =    low_res_image.size[0] * 4
      assert image.shape == (1, expected_height_width, expected_height_width, 3)
      _lowerCAmelCase   :					str        =    np.array([0.3113, 0.3910, 0.4272, 0.4859, 0.5061, 0.4652, 0.5362, 0.5715, 0.5661]       )

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

 def    __magic_name__       (		self :	int       ):






      '''simple docstring'''




      _lowerCAmelCase   :					List[str]        =    "cpu"  # ensure determinism for the device-dependent torch.Generator
      _lowerCAmelCase   :					List[str]        =    self.dummy_cond_unet_upscale
      _lowerCAmelCase   :					str        =    DDPMScheduler()
      _lowerCAmelCase   :					Union[str, Any]        =    DDIMScheduler(prediction_type="v_prediction"       )
      _lowerCAmelCase   :					Tuple        =    self.dummy_vae
      _lowerCAmelCase   :					Union[str, Any]        =    self.dummy_text_encoder
      _lowerCAmelCase   :					List[Any]        =    CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

      _lowerCAmelCase   :					Dict        =    self.dummy_image.cpu().permute(0     ,						2     ,						3     ,						1       )[0]
      _lowerCAmelCase   :					Tuple        =    Image.fromarray(np.uinta(A_       )       ).convert("RGB"       ).resize((6_4, 6_4)       )

      # make sure here that pndm scheduler skips prk
      _lowerCAmelCase   :					Optional[Any]        =    StableDiffusionUpscalePipeline(
          unet=A_     ,						low_res_scheduler=A_     ,						scheduler=A_     ,						vae=A_     ,						text_encoder=A_     ,						tokenizer=A_     ,						max_noise_level=3_5_0     ,						)
      _lowerCAmelCase   :					Dict        =    sd_pipe.to(A_       )
      sd_pipe.set_progress_bar_config(disable=A_       )

      _lowerCAmelCase   :					Union[str, Any]        =    "A painting of a squirrel eating a burger"
      _lowerCAmelCase   :					Tuple        =    sd_pipe(
          2 * [prompt]     ,						image=2 * [low_res_image]     ,						guidance_scale=6.0     ,						noise_level=2_0     ,						num_inference_steps=2     ,						output_type="np"     ,						)
      _lowerCAmelCase   :					List[str]        =    output.images
      assert image.shape[0] == 2

      _lowerCAmelCase   :					str        =    torch.Generator(device=A_       ).manual_seed(0       )
      _lowerCAmelCase   :					List[Any]        =    sd_pipe(
          [prompt]     ,						image=A_     ,						generator=A_     ,						num_images_per_prompt=2     ,						guidance_scale=6.0     ,						noise_level=2_0     ,						num_inference_steps=2     ,						output_type="np"     ,						)
      _lowerCAmelCase   :					int        =    output.images
      assert image.shape[0] == 2


 @unittest.skipIf(torch_device != "cuda"     ,						"This test requires a GPU"       )
 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Any        =    self.dummy_cond_unet_upscale
      _lowerCAmelCase   :					List[str]        =    DDPMScheduler()
      _lowerCAmelCase   :					Optional[int]        =    DDIMScheduler(prediction_type="v_prediction"       )
      _lowerCAmelCase   :					Optional[Any]        =    self.dummy_vae
      _lowerCAmelCase   :					List[str]        =    self.dummy_text_encoder
      _lowerCAmelCase   :					Union[str, Any]        =    CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

      _lowerCAmelCase   :					Any        =    self.dummy_image.cpu().permute(0     ,						2     ,						3     ,						1       )[0]
      _lowerCAmelCase   :					Any        =    Image.fromarray(np.uinta(A_       )       ).convert("RGB"       ).resize((6_4, 6_4)       )

      # put models in fp16, except vae as it overflows in fp16
      _lowerCAmelCase   :					Tuple        =    unet.half()
      _lowerCAmelCase   :					Optional[Any]        =    text_encoder.half()

      # make sure here that pndm scheduler skips prk
      _lowerCAmelCase   :					List[str]        =    StableDiffusionUpscalePipeline(
          unet=A_     ,						low_res_scheduler=A_     ,						scheduler=A_     ,						vae=A_     ,						text_encoder=A_     ,						tokenizer=A_     ,						max_noise_level=3_5_0     ,						)
      _lowerCAmelCase   :					Optional[Any]        =    sd_pipe.to(A_       )
      sd_pipe.set_progress_bar_config(disable=A_       )

      _lowerCAmelCase   :					str        =    "A painting of a squirrel eating a burger"
      _lowerCAmelCase   :					Any        =    torch.manual_seed(0       )
      _lowerCAmelCase   :					Optional[Any]        =    sd_pipe(
          [prompt]     ,						image=A_     ,						generator=A_     ,						num_inference_steps=2     ,						output_type="np"     ,						).images

      _lowerCAmelCase   :					List[str]        =    low_res_image.size[0] * 4
      assert image.shape == (1, expected_height_width, expected_height_width, 3)





@slow
@require_torch_gpu
class A__    ( unittest.TestCase ):



 """simple docstring"""







 def    __magic_name__       (		self :	Any       ):






      '''simple docstring'''




      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()

 def    __magic_name__       (		self :	List[str]       ):






      '''simple docstring'''




      _lowerCAmelCase   :					List[Any]        =    load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
          "/sd2-upscale/low_res_cat.png"       )
      _lowerCAmelCase   :					Dict        =    load_numpy(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
          "/upsampled_cat.npy"       )

      _lowerCAmelCase   :					Union[str, Any]        =    "stabilityai/stable-diffusion-x4-upscaler"
      _lowerCAmelCase   :					Optional[int]        =    StableDiffusionUpscalePipeline.from_pretrained(A_       )
      pipe.to(A_       )
      pipe.set_progress_bar_config(disable=A_       )
      pipe.enable_attention_slicing()

      _lowerCAmelCase   :					str        =    "a cat sitting on a park bench"

      _lowerCAmelCase   :					str        =    torch.manual_seed(0       )
      _lowerCAmelCase   :					Optional[int]        =    pipe(
          prompt=A_     ,						image=A_     ,						generator=A_     ,						output_type="np"     ,						)
      _lowerCAmelCase   :					int        =    output.images[0]

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

 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Optional[int]        =    load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
          "/sd2-upscale/low_res_cat.png"       )
      _lowerCAmelCase   :					Dict        =    load_numpy(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
          "/upsampled_cat_fp16.npy"       )

      _lowerCAmelCase   :					Union[str, Any]        =    "stabilityai/stable-diffusion-x4-upscaler"
      _lowerCAmelCase   :					Any        =    StableDiffusionUpscalePipeline.from_pretrained(
          A_     ,						torch_dtype=torch.floataa     ,						)
      pipe.to(A_       )
      pipe.set_progress_bar_config(disable=A_       )
      pipe.enable_attention_slicing()

      _lowerCAmelCase   :					Optional[Any]        =    "a cat sitting on a park bench"

      _lowerCAmelCase   :					Any        =    torch.manual_seed(0       )
      _lowerCAmelCase   :					int        =    pipe(
          prompt=A_     ,						image=A_     ,						generator=A_     ,						output_type="np"     ,						)
      _lowerCAmelCase   :					Union[str, Any]        =    output.images[0]

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


 def    __magic_name__       (		self :	List[Any]       ):






      '''simple docstring'''




      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

      _lowerCAmelCase   :					int        =    load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
          "/sd2-upscale/low_res_cat.png"       )

      _lowerCAmelCase   :					Tuple        =    "stabilityai/stable-diffusion-x4-upscaler"
      _lowerCAmelCase   :					Any        =    StableDiffusionUpscalePipeline.from_pretrained(
          A_     ,						torch_dtype=torch.floataa     ,						)
      pipe.to(A_       )
      pipe.set_progress_bar_config(disable=A_       )
      pipe.enable_attention_slicing(1       )
      pipe.enable_sequential_cpu_offload()

      _lowerCAmelCase   :					List[Any]        =    "a cat sitting on a park bench"

      _lowerCAmelCase   :					Union[str, Any]        =    torch.manual_seed(0       )
      _lowerCAmelCase   :					List[Any]        =    pipe(
          prompt=A_     ,						image=A_     ,						generator=A_     ,						num_inference_steps=5     ,						output_type="np"     ,						)

      _lowerCAmelCase   :					Optional[Any]        =    torch.cuda.max_memory_allocated()
      # make sure that less than 2.9 GB is allocated
      assert mem_bytes < 2.9 * 1_0**9






