




'''simple docstring'''



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 						lowercase_       (					unittest.TestCase  ):


   """simple docstring"""

   def 				lowerCAmelCase_     (       self   :      Optional[Any]       ):

      """simple docstring"""
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()



   @property
   def 				lowerCAmelCase_     (       self   :      Dict       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= 1
      _SCREAMING_SNAKE_CASE			= 3
      _SCREAMING_SNAKE_CASE			= (3_2, 3_2)

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



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

      """simple docstring"""
      torch.manual_seed(0       )
      _SCREAMING_SNAKE_CASE			= 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=__lowerCamelCase    , only_cross_attention=(True, True, False)    , num_class_embeds=1_0_0    , )
      return model



   @property
   def 				lowerCAmelCase_     (       self   :      Dict       ):

      """simple docstring"""
      torch.manual_seed(0       )
      _SCREAMING_SNAKE_CASE			= 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 				lowerCAmelCase_     (       self   :      List[str]       ):

      """simple docstring"""
      torch.manual_seed(0       )
      _SCREAMING_SNAKE_CASE			= 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(__lowerCamelCase       )



   def 				lowerCAmelCase_     (       self   :      List[str]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= '''cpu'''  # ensure determinism for the device-dependent torch.Generator
      _SCREAMING_SNAKE_CASE			= self.dummy_cond_unet_upscale
      _SCREAMING_SNAKE_CASE			= DDPMScheduler()
      _SCREAMING_SNAKE_CASE			= DDIMScheduler(prediction_type="v_prediction"       )
      _SCREAMING_SNAKE_CASE			= self.dummy_vae
      _SCREAMING_SNAKE_CASE			= self.dummy_text_encoder
      _SCREAMING_SNAKE_CASE			= CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

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

      # make sure here that pndm scheduler skips prk
      _SCREAMING_SNAKE_CASE			= StableDiffusionUpscalePipeline(
          unet=__lowerCamelCase    , low_res_scheduler=__lowerCamelCase    , scheduler=__lowerCamelCase    , vae=__lowerCamelCase    , text_encoder=__lowerCamelCase    , tokenizer=__lowerCamelCase    , max_noise_level=3_5_0    , )
      _SCREAMING_SNAKE_CASE			= sd_pipe.to(__lowerCamelCase       )
      sd_pipe.set_progress_bar_config(disable=__lowerCamelCase       )

      _SCREAMING_SNAKE_CASE			= '''A painting of a squirrel eating a burger'''
      _SCREAMING_SNAKE_CASE			= torch.Generator(device=__lowerCamelCase       ).manual_seed(0       )
      _SCREAMING_SNAKE_CASE			= sd_pipe(
          [prompt]    , image=__lowerCamelCase    , generator=__lowerCamelCase    , guidance_scale=6.0    , noise_level=2_0    , num_inference_steps=2    , output_type="np"    , )

      _SCREAMING_SNAKE_CASE			= output.images

      _SCREAMING_SNAKE_CASE			= torch.Generator(device=__lowerCamelCase       ).manual_seed(0       )
      _SCREAMING_SNAKE_CASE			= sd_pipe(
          [prompt]    , image=__lowerCamelCase    , generator=__lowerCamelCase    , guidance_scale=6.0    , noise_level=2_0    , num_inference_steps=2    , output_type="np"    , return_dict=__lowerCamelCase    , )[0]

      _SCREAMING_SNAKE_CASE			= image[0, -3:, -3:, -1]
      _SCREAMING_SNAKE_CASE			= image_from_tuple[0, -3:, -3:, -1]

      _SCREAMING_SNAKE_CASE			= low_res_image.size[0] * 4
      assert image.shape == (1, expected_height_width, expected_height_width, 3)
      _SCREAMING_SNAKE_CASE			= np.array([0.3_1_1_3, 0.3_9_1_0, 0.4_2_7_2, 0.4_8_5_9, 0.5_0_6_1, 0.4_6_5_2, 0.5_3_6_2, 0.5_7_1_5, 0.5_6_6_1]       )

      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 				lowerCAmelCase_     (       self   :      List[str]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= '''cpu'''  # ensure determinism for the device-dependent torch.Generator
      _SCREAMING_SNAKE_CASE			= self.dummy_cond_unet_upscale
      _SCREAMING_SNAKE_CASE			= DDPMScheduler()
      _SCREAMING_SNAKE_CASE			= DDIMScheduler(prediction_type="v_prediction"       )
      _SCREAMING_SNAKE_CASE			= self.dummy_vae
      _SCREAMING_SNAKE_CASE			= self.dummy_text_encoder
      _SCREAMING_SNAKE_CASE			= CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

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

      # make sure here that pndm scheduler skips prk
      _SCREAMING_SNAKE_CASE			= StableDiffusionUpscalePipeline(
          unet=__lowerCamelCase    , low_res_scheduler=__lowerCamelCase    , scheduler=__lowerCamelCase    , vae=__lowerCamelCase    , text_encoder=__lowerCamelCase    , tokenizer=__lowerCamelCase    , max_noise_level=3_5_0    , )
      _SCREAMING_SNAKE_CASE			= sd_pipe.to(__lowerCamelCase       )
      sd_pipe.set_progress_bar_config(disable=__lowerCamelCase       )

      _SCREAMING_SNAKE_CASE			= '''A painting of a squirrel eating a burger'''
      _SCREAMING_SNAKE_CASE			= sd_pipe(
          2 * [prompt]    , image=2 * [low_res_image]    , guidance_scale=6.0    , noise_level=2_0    , num_inference_steps=2    , output_type="np"    , )
      _SCREAMING_SNAKE_CASE			= output.images
      assert image.shape[0] == 2

      _SCREAMING_SNAKE_CASE			= torch.Generator(device=__lowerCamelCase       ).manual_seed(0       )
      _SCREAMING_SNAKE_CASE			= sd_pipe(
          [prompt]    , image=__lowerCamelCase    , generator=__lowerCamelCase    , num_images_per_prompt=2    , guidance_scale=6.0    , noise_level=2_0    , num_inference_steps=2    , output_type="np"    , )
      _SCREAMING_SNAKE_CASE			= output.images
      assert image.shape[0] == 2



   @unittest.skipIf(torch_device != "cuda"    , "This test requires a GPU"       )
   def 				lowerCAmelCase_     (       self   :      Union[str, Any]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= self.dummy_cond_unet_upscale
      _SCREAMING_SNAKE_CASE			= DDPMScheduler()
      _SCREAMING_SNAKE_CASE			= DDIMScheduler(prediction_type="v_prediction"       )
      _SCREAMING_SNAKE_CASE			= self.dummy_vae
      _SCREAMING_SNAKE_CASE			= self.dummy_text_encoder
      _SCREAMING_SNAKE_CASE			= CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

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

      # put models in fp16, except vae as it overflows in fp16
      _SCREAMING_SNAKE_CASE			= unet.half()
      _SCREAMING_SNAKE_CASE			= text_encoder.half()

      # make sure here that pndm scheduler skips prk
      _SCREAMING_SNAKE_CASE			= StableDiffusionUpscalePipeline(
          unet=__lowerCamelCase    , low_res_scheduler=__lowerCamelCase    , scheduler=__lowerCamelCase    , vae=__lowerCamelCase    , text_encoder=__lowerCamelCase    , tokenizer=__lowerCamelCase    , max_noise_level=3_5_0    , )
      _SCREAMING_SNAKE_CASE			= sd_pipe.to(__lowerCamelCase       )
      sd_pipe.set_progress_bar_config(disable=__lowerCamelCase       )

      _SCREAMING_SNAKE_CASE			= '''A painting of a squirrel eating a burger'''
      _SCREAMING_SNAKE_CASE			= torch.manual_seed(0       )
      _SCREAMING_SNAKE_CASE			= sd_pipe(
          [prompt]    , image=__lowerCamelCase    , generator=__lowerCamelCase    , num_inference_steps=2    , output_type="np"    , ).images

      _SCREAMING_SNAKE_CASE			= low_res_image.size[0] * 4
      assert image.shape == (1, expected_height_width, expected_height_width, 3)







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


   """simple docstring"""

   def 				lowerCAmelCase_     (       self   :      str       ):

      """simple docstring"""
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()



   def 				lowerCAmelCase_     (       self   :      int       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
          "/sd2-upscale/low_res_cat.png"       )
      _SCREAMING_SNAKE_CASE			= load_numpy(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
          "/upsampled_cat.npy"       )

      _SCREAMING_SNAKE_CASE			= '''stabilityai/stable-diffusion-x4-upscaler'''
      _SCREAMING_SNAKE_CASE			= StableDiffusionUpscalePipeline.from_pretrained(__lowerCamelCase       )
      pipe.to(__lowerCamelCase       )
      pipe.set_progress_bar_config(disable=__lowerCamelCase       )
      pipe.enable_attention_slicing()

      _SCREAMING_SNAKE_CASE			= '''a cat sitting on a park bench'''

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

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



   def 				lowerCAmelCase_     (       self   :      Optional[int]       ):

      """simple docstring"""
      _SCREAMING_SNAKE_CASE			= load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
          "/sd2-upscale/low_res_cat.png"       )
      _SCREAMING_SNAKE_CASE			= load_numpy(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
          "/upsampled_cat_fp16.npy"       )

      _SCREAMING_SNAKE_CASE			= '''stabilityai/stable-diffusion-x4-upscaler'''
      _SCREAMING_SNAKE_CASE			= StableDiffusionUpscalePipeline.from_pretrained(
          __lowerCamelCase    , torch_dtype=torch.floataa    , )
      pipe.to(__lowerCamelCase       )
      pipe.set_progress_bar_config(disable=__lowerCamelCase       )
      pipe.enable_attention_slicing()

      _SCREAMING_SNAKE_CASE			= '''a cat sitting on a park bench'''

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

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



   def 				lowerCAmelCase_     (       self   :      Any       ):

      """simple docstring"""
      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

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

      _SCREAMING_SNAKE_CASE			= '''stabilityai/stable-diffusion-x4-upscaler'''
      _SCREAMING_SNAKE_CASE			= StableDiffusionUpscalePipeline.from_pretrained(
          __lowerCamelCase    , torch_dtype=torch.floataa    , )
      pipe.to(__lowerCamelCase       )
      pipe.set_progress_bar_config(disable=__lowerCamelCase       )
      pipe.enable_attention_slicing(1       )
      pipe.enable_sequential_cpu_offload()

      _SCREAMING_SNAKE_CASE			= '''a cat sitting on a park bench'''

      _SCREAMING_SNAKE_CASE			= torch.manual_seed(0       )
      _SCREAMING_SNAKE_CASE			= pipe(
          prompt=__lowerCamelCase    , image=__lowerCamelCase    , generator=__lowerCamelCase    , num_inference_steps=5    , output_type="np"    , )

      _SCREAMING_SNAKE_CASE			= torch.cuda.max_memory_allocated()
      # make sure that less than 2.9 GB is allocated
      assert mem_bytes < 2.9 * 1_0**9








"""simple docstring"""



import os
import textwrap

import pyarrow as pa
import pytest

from datasets import ClassLabel, Features, Image
from datasets.packaged_modules.csv.csv import Csv

from ..utils import require_pil







@pytest.fixture
def 						__lowercase	(  snake_case_    :    Optional[int]		)  ->str:



    '''simple docstring'''







    __A       :		Tuple							=	tmp_path / '''file.csv'''
    __A       :		Union[str, Any]							=	textwrap.dedent(
        '''\
        header1,header2
        1,2
        10,20
        '''		)
    with open(snake_case_							,'''w'''		) as f:
        f.write(snake_case_		)
    return str(snake_case_		)







@pytest.fixture
def 						__lowercase	(  snake_case_    :    str		)  ->Optional[Any]:



    '''simple docstring'''







    __A       :		List[str]							=	tmp_path / '''malformed_file.csv'''
    __A       :		Optional[Any]							=	textwrap.dedent(
        '''\
        header1,header2
        1,2
        10,20,
        '''		)
    with open(snake_case_							,'''w'''		) as f:
        f.write(snake_case_		)
    return str(snake_case_		)







@pytest.fixture
def 						__lowercase	(  snake_case_    :    Any							,snake_case_    :    Optional[int]		)  ->Optional[Any]:



    '''simple docstring'''







    __A       :		List[str]							=	tmp_path / '''csv_with_image.csv'''
    __A       :		List[str]							=	textwrap.dedent(
        F"""\
        image
        {image_file}
        """		)
    with open(snake_case_							,'''w'''		) as f:
        f.write(snake_case_		)
    return str(snake_case_		)







@pytest.fixture
def 						__lowercase	(  snake_case_    :    Optional[int]		)  ->Tuple:



    '''simple docstring'''







    __A       :		Dict							=	tmp_path / '''csv_with_label.csv'''
    __A       :		int							=	textwrap.dedent(
        '''\
        label
        good
        bad
        good
        '''		)
    with open(snake_case_							,'''w'''		) as f:
        f.write(snake_case_		)
    return str(snake_case_		)







@pytest.fixture
def 						__lowercase	(  snake_case_    :    Dict		)  ->Optional[Any]:



    '''simple docstring'''







    __A       :		Optional[int]							=	tmp_path / '''csv_with_int_list.csv'''
    __A       :		Union[str, Any]							=	textwrap.dedent(
        '''\
        int_list
        1 2 3
        4 5 6
        7 8 9
        '''		)
    with open(snake_case_							,'''w'''		) as f:
        f.write(snake_case_		)
    return str(snake_case_		)







def 						__lowercase	(  snake_case_    :    List[str]							,snake_case_    :    Optional[int]							,snake_case_    :    Optional[Any]		)  ->Dict:



    '''simple docstring'''







    __A       :		str							=	Csv()
    __A       :		List[Any]							=	csv._generate_tables([[csv_file, malformed_csv_file]]		)
    with pytest.raises(snake_case_							,match='''Error tokenizing data'''		):
        for _ in generator:
            pass
    assert any(
        record.levelname == '''ERROR'''
        and '''Failed to read file''' in record.message
        and os.path.basename(snake_case_		) in record.message
        for record in caplog.records		)







@require_pil
def 						__lowercase	(  snake_case_    :    Optional[int]		)  ->Union[str, Any]:



    '''simple docstring'''







    with open(snake_case_							,encoding='''utf-8'''		) as f:
        __A       :		int							=	f.read().splitlines()[1]
    __A       :		Optional[Any]							=	Csv(encoding='''utf-8'''							,features=Features({'''image''': Image()}		)		)
    __A       :		Union[str, Any]							=	csv._generate_tables([[csv_file_with_image]]		)
    __A       :		Optional[Any]							=	pa.concat_tables([table for _, table in generator]		)
    assert pa_table.schema.field('''image'''		).type == Image()()
    __A       :		Optional[int]							=	pa_table.to_pydict()['''image''']
    assert generated_content == [{"path": image_file, "bytes": None}]







def 						__lowercase	(  snake_case_    :    Any		)  ->List[Any]:



    '''simple docstring'''







    with open(snake_case_							,encoding='''utf-8'''		) as f:
        __A       :		Optional[int]							=	f.read().splitlines()[1:]
    __A       :		List[Any]							=	Csv(encoding='''utf-8'''							,features=Features({'''label''': ClassLabel(names=['''good''', '''bad''']		)}		)		)
    __A       :		str							=	csv._generate_tables([[csv_file_with_label]]		)
    __A       :		Any							=	pa.concat_tables([table for _, table in generator]		)
    assert pa_table.schema.field('''label'''		).type == ClassLabel(names=['''good''', '''bad''']		)()
    __A       :		Optional[Any]							=	pa_table.to_pydict()['''label''']
    assert generated_content == [ClassLabel(names=['''good''', '''bad''']		).straint(snake_case_		) for label in labels]







def 						__lowercase	(  snake_case_    :    str		)  ->Dict:



    '''simple docstring'''







    __A       :		Optional[int]							=	Csv(encoding='''utf-8'''							,sep=''','''							,converters={'''int_list''': lambda snake_case_		: [int(snake_case_		) for i in x.split()]}		)
    __A       :		Dict							=	csv._generate_tables([[csv_file_with_int_list]]		)
    __A       :		Dict							=	pa.concat_tables([table for _, table in generator]		)
    assert pa.types.is_list(pa_table.schema.field('''int_list'''		).type		)
    __A       :		int							=	pa_table.to_pydict()['''int_list''']
    assert generated_content == [[1, 2, 3], [4, 5, 6], [7, 8, 9]]




