





'''simple docstring'''
import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, LDMSuperResolutionPipeline, UNetaDModel, VQModel
from diffusers.utils import PIL_INTERPOLATION, floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch


enable_full_determinism()

class 					snake_case					(   unittest.TestCase	):




 """simple docstring"""



 @property
 def 	_lowerCamelCase (	self	:							Dict     ):
        __UpperCamelCase	=  1
        __UpperCamelCase	=  3
        __UpperCamelCase	=  (3_2, 3_2)

        __UpperCamelCase	=  floats_tensor((batch_size, num_channels) + sizes							,     rng=random.Random(0     )     ).to(__A     )
        return image






 @property
 def 	_lowerCamelCase (	self	:							Optional[int]     ):
        torch.manual_seed(0     )
        __UpperCamelCase	=  UNetaDModel(
            block_out_channels=(3_2, 6_4)							,     layers_per_block=2							,     sample_size=3_2							,     in_channels=6							,     out_channels=3							,     down_block_types=('DownBlock2D', 'AttnDownBlock2D')							,     up_block_types=('AttnUpBlock2D', 'UpBlock2D')							,     )
        return model






 @property
 def 	_lowerCamelCase (	self	:							Dict     ):
        torch.manual_seed(0     )
        __UpperCamelCase	=  VQModel(
            block_out_channels=[3_2, 6_4]							,     in_channels=3							,     out_channels=3							,     down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']							,     up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']							,     latent_channels=3							,     )
        return model






 def 	_lowerCamelCase (	self	:							Any     ):
        __UpperCamelCase	=  'cpu'
        __UpperCamelCase	=  self.dummy_uncond_unet
        __UpperCamelCase	=  DDIMScheduler()
        __UpperCamelCase	=  self.dummy_vq_model

        __UpperCamelCase	=  LDMSuperResolutionPipeline(unet=__A							,     vqvae=__A							,     scheduler=__A     )
        ldm.to(__A     )
        ldm.set_progress_bar_config(disable=__A     )

        __UpperCamelCase	=  self.dummy_image.to(__A     )

        __UpperCamelCase	=  torch.Generator(device=__A     ).manual_seed(0     )
        __UpperCamelCase	=  ldm(image=__A							,     generator=__A							,     num_inference_steps=2							,     output_type='numpy'     ).images

        __UpperCamelCase	=  image[0, -3:, -3:, -1]

        assert image.shape == (1, 6_4, 6_4, 3)
        __UpperCamelCase	=  np.array([0.8678, 0.8245, 0.6381, 0.6830, 0.4385, 0.5599, 0.4641, 0.6201, 0.5150]     )

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






 @unittest.skipIf(torch_device != 'cuda'							,     'This test requires a GPU'     )
 def 	_lowerCamelCase (	self	:							Optional[int]     ):
        __UpperCamelCase	=  self.dummy_uncond_unet
        __UpperCamelCase	=  DDIMScheduler()
        __UpperCamelCase	=  self.dummy_vq_model

        # put models in fp16
        __UpperCamelCase	=  unet.half()
        __UpperCamelCase	=  vqvae.half()

        __UpperCamelCase	=  LDMSuperResolutionPipeline(unet=__A							,     vqvae=__A							,     scheduler=__A     )
        ldm.to(__A     )
        ldm.set_progress_bar_config(disable=__A     )

        __UpperCamelCase	=  self.dummy_image.to(__A     )

        __UpperCamelCase	=  ldm(__A							,     num_inference_steps=2							,     output_type='numpy'     ).images

        assert image.shape == (1, 6_4, 6_4, 3)






@slow
@require_torch
class 					snake_case					(   unittest.TestCase	):




 """simple docstring"""



 def 	_lowerCamelCase (	self	:							Dict     ):
        __UpperCamelCase	=  load_image(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/vq_diffusion/teddy_bear_pool.png'     )
        __UpperCamelCase	=  init_image.resize((6_4, 6_4)							,     resample=PIL_INTERPOLATION['lanczos']     )

        __UpperCamelCase	=  LDMSuperResolutionPipeline.from_pretrained('duongna/ldm-super-resolution'							,     device_map='auto'     )
        ldm.set_progress_bar_config(disable=__A     )

        __UpperCamelCase	=  torch.manual_seed(0     )
        __UpperCamelCase	=  ldm(image=__A							,     generator=__A							,     num_inference_steps=2_0							,     output_type='numpy'     ).images

        __UpperCamelCase	=  image[0, -3:, -3:, -1]

        assert image.shape == (1, 2_5_6, 2_5_6, 3)
        __UpperCamelCase	=  np.array([0.7644, 0.7679, 0.7642, 0.7633, 0.7666, 0.7560, 0.7425, 0.7257, 0.6907]     )

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







'''simple docstring'''
import itertools
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa

import datasets
from datasets.table import table_cast


a__							:							Optional[Any]			=datasets.utils.logging.get_logger(__name__)

@dataclass
class 					snake_case					(   datasets.BuilderConfig	):




 """simple docstring"""



 SCREAMING_SNAKE_CASE_     : Optional[datasets.Features]      		=None






class 					snake_case					(   datasets.ArrowBasedBuilder	):




 """simple docstring"""



 SCREAMING_SNAKE_CASE_     : int      		=ArrowConfig






 def 	_lowerCamelCase (	self	:							List[str]     ):
        return datasets.DatasetInfo(features=self.config.features     )






 def 	_lowerCamelCase (	self	:							Union[str, Any]							,     __A	:							Dict     ):
        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}'''     )
        __UpperCamelCase	=  dl_manager.download_and_extract(self.config.data_files     )
        if isinstance(__A							,     (str, list, tuple)     ):
               __UpperCamelCase	=  data_files
               if isinstance(__A							,     __A     ):
                      __UpperCamelCase	=  [files]
               # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
               __UpperCamelCase	=  [dl_manager.iter_files(__A     ) for file in files]
               return [datasets.SplitGenerator(name=datasets.Split.TRAIN							,     gen_kwargs={'files': files}     )]
        __UpperCamelCase	=  []
        for split_name, files in data_files.items():
               if isinstance(__A							,     __A     ):
                      __UpperCamelCase	=  [files]
               # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
               __UpperCamelCase	=  [dl_manager.iter_files(__A     ) for file in files]
               # Infer features is they are stoed in the arrow schema
               if self.info.features is None:
                      for file in itertools.chain.from_iterable(__A     ):
                             with open(__A							,     'rb'     ) as f:
                                    __UpperCamelCase	=  datasets.Features.from_arrow_schema(pa.ipc.open_stream(__A     ).schema     )
                             break
               splits.append(datasets.SplitGenerator(name=__A							,     gen_kwargs={'files': files}     )     )
        return splits






 def 	_lowerCamelCase (	self	:							Any							,     __A	:							pa.Table     ):
        if self.info.features is not None:
               # more expensive cast to support nested features with keys in a different order
               # allows str <-> int/float or str to Audio for example
               __UpperCamelCase	=  table_cast(__A							,     self.info.features.arrow_schema     )
        return pa_table






 def 	_lowerCamelCase (	self	:							int							,     __A	:							Optional[Any]     ):
        for file_idx, file in enumerate(itertools.chain.from_iterable(__A     )     ):
               with open(__A							,     'rb'     ) as f:
                      try:
                             for batch_idx, record_batch in enumerate(pa.ipc.open_stream(__A     )     ):
                                    __UpperCamelCase	=  pa.Table.from_batches([record_batch]     )
                                    # 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 f'''{file_idx}_{batch_idx}''', self._cast_table(__A     )
                      except ValueError as e:
                             logger.error(f'''Failed to read file \'{file}\' with error {type(__A     )}: {e}'''     )
                             raise
