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

      def 				A__       (						self						):







            """simple docstring"""



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

      @property
      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	1
            lowercase														=	3
            lowercase														=	(32, 32)

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

      @property
      def 				A__       (						self						):







            """simple docstring"""



            torch.manual_seed(0						)
            lowercase														=	UNetaDConditionModel(
                block_out_channels=(32, 32, 64)		,       layers_per_block=2		,       sample_size=32		,       in_channels=7		,       out_channels=4		,       down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""", """CrossAttnDownBlock2D""")		,       up_block_types=("""CrossAttnUpBlock2D""", """CrossAttnUpBlock2D""", """UpBlock2D""")		,       cross_attention_dim=32		,       attention_head_dim=8		,       use_linear_projection=_A		,       only_cross_attention=(True, True, False)		,       num_class_embeds=100		,       )
            return model

      @property
      def 				A__       (						self						):







            """simple docstring"""



            torch.manual_seed(0						)
            lowercase														=	AutoencoderKL(
                block_out_channels=[32, 32, 64]		,       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 				A__       (						self						):







            """simple docstring"""



            torch.manual_seed(0						)
            lowercase														=	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=1000		,       hidden_act="""gelu"""		,       projection_dim=512		,       )
            return CLIPTextModel(_A						)

      def 				A__       (						self						):







            """simple docstring"""



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

            lowercase														=	self.dummy_image.cpu().permute(0		,       2		,       3		,       1						)[0]
            lowercase														=	Image.fromarray(np.uinta(_A						)						).convert("""RGB"""						).resize((64, 64)						)

            # make sure here that pndm scheduler skips prk
            lowercase														=	StableDiffusionUpscalePipeline(
                unet=_A		,       low_res_scheduler=_A		,       scheduler=_A		,       vae=_A		,       text_encoder=_A		,       tokenizer=_A		,       max_noise_level=350		,       )
            lowercase														=	sd_pipe.to(_A						)
            sd_pipe.set_progress_bar_config(disable=_A						)

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

            lowercase														=	output.images

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

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

            lowercase														=	low_res_image.size[0] * 4
            assert image.shape == (1, expected_height_width, expected_height_width, 3)
            lowercase														=	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 				A__       (						self						):







            """simple docstring"""



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

            lowercase														=	self.dummy_image.cpu().permute(0		,       2		,       3		,       1						)[0]
            lowercase														=	Image.fromarray(np.uinta(_A						)						).convert("""RGB"""						).resize((64, 64)						)

            # make sure here that pndm scheduler skips prk
            lowercase														=	StableDiffusionUpscalePipeline(
                unet=_A		,       low_res_scheduler=_A		,       scheduler=_A		,       vae=_A		,       text_encoder=_A		,       tokenizer=_A		,       max_noise_level=350		,       )
            lowercase														=	sd_pipe.to(_A						)
            sd_pipe.set_progress_bar_config(disable=_A						)

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

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







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







            """simple docstring"""



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

            lowercase														=	self.dummy_image.cpu().permute(0		,       2		,       3		,       1						)[0]
            lowercase														=	Image.fromarray(np.uinta(_A						)						).convert("""RGB"""						).resize((64, 64)						)

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

            # make sure here that pndm scheduler skips prk
            lowercase														=	StableDiffusionUpscalePipeline(
                unet=_A		,       low_res_scheduler=_A		,       scheduler=_A		,       vae=_A		,       text_encoder=_A		,       tokenizer=_A		,       max_noise_level=350		,       )
            lowercase														=	sd_pipe.to(_A						)
            sd_pipe.set_progress_bar_config(disable=_A						)

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

            lowercase														=	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	):

      def 				A__       (						self						):







            """simple docstring"""



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

      def 				A__       (						self						):







            """simple docstring"""



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

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

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

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

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

      def 				A__       (						self						):







            """simple docstring"""



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

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

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

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

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







      def 				A__       (						self						):







            """simple docstring"""



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

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

            lowercase														=	'''stabilityai/stable-diffusion-x4-upscaler'''
            lowercase														=	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()

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

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

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







'''simple docstring'''
import os
from collections import namedtuple

import pytest

from datasets import ClassLabel, Features, Sequence, Value
from datasets.commands.test import TestCommand
from datasets.info import DatasetInfo, DatasetInfosDict


UpperCamelCase__    					=      namedtuple(
    '''_TestCommandArgs''',
    [
        '''dataset''',
        '''name''',
        '''cache_dir''',
        '''data_dir''',
        '''all_configs''',
        '''save_infos''',
        '''ignore_verifications''',
        '''force_redownload''',
        '''clear_cache''',
    ],
    defaults=[None, None, None, False, False, False, False, False],
)





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	Optional[int]:
    return (abs(source - target						) / target) < 0.0_1







@pytest.mark.integration
def 					a__       (					lowerCAmelCase__						)    ->	Optional[int]:
    UpperCAmelCase__			:		str													=				_TestCommandArgs(dataset=lowerCAmelCase__  ,					all_configs=lowerCAmelCase__  ,					save_infos=lowerCAmelCase__						)
    UpperCAmelCase__			:		Tuple													=				TestCommand(*lowerCAmelCase__						)
    test_command.run()
    UpperCAmelCase__			:		List[Any]													=				os.path.join(lowerCAmelCase__  ,					'''README.md'''						)
    assert os.path.exists(lowerCAmelCase__						)
    UpperCAmelCase__			:		List[str]													=				DatasetInfosDict.from_directory(lowerCAmelCase__						)
    UpperCAmelCase__			:		Any													=				DatasetInfosDict(
        {
            '''default''': DatasetInfo(
                features=Features(
                    {
                        '''tokens''': Sequence(Value('''string'''						)						),
                        '''ner_tags''': Sequence(
                            ClassLabel(names=['''O''', '''B-PER''', '''I-PER''', '''B-ORG''', '''I-ORG''', '''B-LOC''', '''I-LOC''']						)						),
                        '''langs''': Sequence(Value('''string'''						)						),
                        '''spans''': Sequence(Value('''string'''						)						),
                    }						)  ,					splits=[
                    {
                        '''name''': '''train''',
                        '''num_bytes''': 2_35_15_63,
                        '''num_examples''': 1_00_00,
                    },
                    {
                        '''name''': '''validation''',
                        '''num_bytes''': 23_84_18,
                        '''num_examples''': 10_00,
                    },
                ]  ,					download_size=3_94_06_80  ,					dataset_size=2_58_99_81  ,					)
        }						)
    assert dataset_infos.keys() == expected_dataset_infos.keys()
    for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
        UpperCAmelCase__						,						UpperCAmelCase__			:		str													=				getattr(dataset_infos['''default''']  ,					lowerCAmelCase__						), getattr(expected_dataset_infos['''default''']  ,					lowerCAmelCase__						)
        if key == "num_bytes":
            assert is_apercent_close(lowerCAmelCase__  ,					lowerCAmelCase__						)
        elif key == "splits":
            assert list(lowerCAmelCase__						) == list(lowerCAmelCase__						)
            for split in result:
                assert result[split].name == expected[split].name
                assert result[split].num_examples == expected[split].num_examples
                assert is_apercent_close(result[split].num_bytes  ,					expected[split].num_bytes						)
        else:
            result == expected
