



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







	'''simple docstring'''






	def  a_					(      self					):
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()






	@property
	def  a_					(      self					):
						__SCREAMING_SNAKE_CASE				:    List[Any]						    =							1
						__SCREAMING_SNAKE_CASE				:    Dict						    =							3
						__SCREAMING_SNAKE_CASE				:    int						    =							(32, 32)

						__SCREAMING_SNAKE_CASE				:    Optional[Any]						    =							floats_tensor((batch_size, num_channels) + sizes      ,     rng=random.Random(0					)					).to(__lowerCamelCase					)
						return image






	@property
	def  a_					(      self					):
						torch.manual_seed(0					)
						__SCREAMING_SNAKE_CASE				:    List[str]						    =							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=__lowerCamelCase      ,     only_cross_attention=(True, True, False)      ,     num_class_embeds=100      ,     )
						return model






	@property
	def  a_					(      self					):
						torch.manual_seed(0					)
						__SCREAMING_SNAKE_CASE				:    Any						    =							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					):
						torch.manual_seed(0					)
						__SCREAMING_SNAKE_CASE				:    Tuple						    =							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(__lowerCamelCase					)






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

						__SCREAMING_SNAKE_CASE				:    List[Any]						    =							self.dummy_image.cpu().permute(0      ,     2      ,     3      ,     1					)[0]
						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							Image.fromarray(np.uinta(__lowerCamelCase					)					).convert("RGB"					).resize((64, 64)					)

						# make sure here that pndm scheduler skips prk
						__SCREAMING_SNAKE_CASE				:    Tuple						    =							StableDiffusionUpscalePipeline(
						    unet=__lowerCamelCase      ,     low_res_scheduler=__lowerCamelCase      ,     scheduler=__lowerCamelCase      ,     vae=__lowerCamelCase      ,     text_encoder=__lowerCamelCase      ,     tokenizer=__lowerCamelCase      ,     max_noise_level=350      ,     )
						__SCREAMING_SNAKE_CASE				:    int						    =							sd_pipe.to(__lowerCamelCase					)
						sd_pipe.set_progress_bar_config(disable=__lowerCamelCase					)

						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							'''A painting of a squirrel eating a burger'''
						__SCREAMING_SNAKE_CASE				:    Tuple						    =							torch.Generator(device=__lowerCamelCase					).manual_seed(0					)
						__SCREAMING_SNAKE_CASE				:    Dict						    =							sd_pipe(
						    [prompt]      ,     image=__lowerCamelCase      ,     generator=__lowerCamelCase      ,     guidance_scale=6.0      ,     noise_level=20      ,     num_inference_steps=2      ,     output_type="np"      ,     )

						__SCREAMING_SNAKE_CASE				:    List[Any]						    =							output.images

						__SCREAMING_SNAKE_CASE				:    Dict						    =							torch.Generator(device=__lowerCamelCase					).manual_seed(0					)
						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							sd_pipe(
						    [prompt]      ,     image=__lowerCamelCase      ,     generator=__lowerCamelCase      ,     guidance_scale=6.0      ,     noise_level=20      ,     num_inference_steps=2      ,     output_type="np"      ,     return_dict=__lowerCamelCase      ,     )[0]

						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							image[0, -3:, -3:, -1]
						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							image_from_tuple[0, -3:, -3:, -1]

						__SCREAMING_SNAKE_CASE				:    Dict						    =							low_res_image.size[0] * 4
						assert image.shape == (1, expected_height_width, expected_height_width, 3)
						__SCREAMING_SNAKE_CASE				:    int						    =							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  a_					(      self					):
						__SCREAMING_SNAKE_CASE				:    List[Any]						    =							'''cpu'''  # ensure determinism for the device-dependent torch.Generator
						__SCREAMING_SNAKE_CASE				:    Any						    =							self.dummy_cond_unet_upscale
						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							DDPMScheduler()
						__SCREAMING_SNAKE_CASE				:    Dict						    =							DDIMScheduler(prediction_type="v_prediction"					)
						__SCREAMING_SNAKE_CASE				:    Optional[int]						    =							self.dummy_vae
						__SCREAMING_SNAKE_CASE				:    Dict						    =							self.dummy_text_encoder
						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					)

						__SCREAMING_SNAKE_CASE				:    List[str]						    =							self.dummy_image.cpu().permute(0      ,     2      ,     3      ,     1					)[0]
						__SCREAMING_SNAKE_CASE				:    Optional[Any]						    =							Image.fromarray(np.uinta(__lowerCamelCase					)					).convert("RGB"					).resize((64, 64)					)

						# make sure here that pndm scheduler skips prk
						__SCREAMING_SNAKE_CASE				:    List[str]						    =							StableDiffusionUpscalePipeline(
						    unet=__lowerCamelCase      ,     low_res_scheduler=__lowerCamelCase      ,     scheduler=__lowerCamelCase      ,     vae=__lowerCamelCase      ,     text_encoder=__lowerCamelCase      ,     tokenizer=__lowerCamelCase      ,     max_noise_level=350      ,     )
						__SCREAMING_SNAKE_CASE				:    Dict						    =							sd_pipe.to(__lowerCamelCase					)
						sd_pipe.set_progress_bar_config(disable=__lowerCamelCase					)

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

						__SCREAMING_SNAKE_CASE				:    Dict						    =							torch.Generator(device=__lowerCamelCase					).manual_seed(0					)
						__SCREAMING_SNAKE_CASE				:    Dict						    =							sd_pipe(
						    [prompt]      ,     image=__lowerCamelCase      ,     generator=__lowerCamelCase      ,     num_images_per_prompt=2      ,     guidance_scale=6.0      ,     noise_level=20      ,     num_inference_steps=2      ,     output_type="np"      ,     )
						__SCREAMING_SNAKE_CASE				:    Optional[Any]						    =							output.images
						assert image.shape[0] == 2






	@unittest.skipIf(torch_device != "cuda"      ,     "This test requires a GPU"					)
	def  a_					(      self					):
						__SCREAMING_SNAKE_CASE				:    Dict						    =							self.dummy_cond_unet_upscale
						__SCREAMING_SNAKE_CASE				:    Dict						    =							DDPMScheduler()
						__SCREAMING_SNAKE_CASE				:    int						    =							DDIMScheduler(prediction_type="v_prediction"					)
						__SCREAMING_SNAKE_CASE				:    Any						    =							self.dummy_vae
						__SCREAMING_SNAKE_CASE				:    Optional[int]						    =							self.dummy_text_encoder
						__SCREAMING_SNAKE_CASE				:    List[Any]						    =							CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					)

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

						# put models in fp16, except vae as it overflows in fp16
						__SCREAMING_SNAKE_CASE				:    str						    =							unet.half()
						__SCREAMING_SNAKE_CASE				:    List[str]						    =							text_encoder.half()

						# make sure here that pndm scheduler skips prk
						__SCREAMING_SNAKE_CASE				:    Tuple						    =							StableDiffusionUpscalePipeline(
						    unet=__lowerCamelCase      ,     low_res_scheduler=__lowerCamelCase      ,     scheduler=__lowerCamelCase      ,     vae=__lowerCamelCase      ,     text_encoder=__lowerCamelCase      ,     tokenizer=__lowerCamelCase      ,     max_noise_level=350      ,     )
						__SCREAMING_SNAKE_CASE				:    Dict						    =							sd_pipe.to(__lowerCamelCase					)
						sd_pipe.set_progress_bar_config(disable=__lowerCamelCase					)

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

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






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







	'''simple docstring'''






	def  a_					(      self					):
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()






	def  a_					(      self					):
						__SCREAMING_SNAKE_CASE				:    Optional[Any]						    =							load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/sd2-upscale/low_res_cat.png"					)
						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
						    "/upsampled_cat.npy"					)

						__SCREAMING_SNAKE_CASE				:    str						    =							'''stabilityai/stable-diffusion-x4-upscaler'''
						__SCREAMING_SNAKE_CASE				:    Union[str, Any]						    =							StableDiffusionUpscalePipeline.from_pretrained(__lowerCamelCase					)
						pipe.to(__lowerCamelCase					)
						pipe.set_progress_bar_config(disable=__lowerCamelCase					)
						pipe.enable_attention_slicing()

						__SCREAMING_SNAKE_CASE				:    Optional[Any]						    =							'''a cat sitting on a park bench'''

						__SCREAMING_SNAKE_CASE				:    Dict						    =							torch.manual_seed(0					)
						__SCREAMING_SNAKE_CASE				:    Any						    =							pipe(
						    prompt=__lowerCamelCase      ,     image=__lowerCamelCase      ,     generator=__lowerCamelCase      ,     output_type="np"      ,     )
						__SCREAMING_SNAKE_CASE				:    List[str]						    =							output.images[0]

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






	def  a_					(      self					):
						__SCREAMING_SNAKE_CASE				:    Optional[Any]						    =							load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/sd2-upscale/low_res_cat.png"					)
						__SCREAMING_SNAKE_CASE				:    str						    =							load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
						    "/upsampled_cat_fp16.npy"					)

						__SCREAMING_SNAKE_CASE				:    List[Any]						    =							'''stabilityai/stable-diffusion-x4-upscaler'''
						__SCREAMING_SNAKE_CASE				:    Dict						    =							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				:    str						    =							'''a cat sitting on a park bench'''

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

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






	def  a_					(      self					):
						torch.cuda.empty_cache()
						torch.cuda.reset_max_memory_allocated()
						torch.cuda.reset_peak_memory_stats()

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

						__SCREAMING_SNAKE_CASE				:    Any						    =							'''stabilityai/stable-diffusion-x4-upscaler'''
						__SCREAMING_SNAKE_CASE				:    List[str]						    =							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				:    Any						    =							'''a cat sitting on a park bench'''

						__SCREAMING_SNAKE_CASE				:    List[Any]						    =							torch.manual_seed(0					)
						__SCREAMING_SNAKE_CASE				:    int						    =							pipe(
						    prompt=__lowerCamelCase      ,     image=__lowerCamelCase      ,     generator=__lowerCamelCase      ,     num_inference_steps=5      ,     output_type="np"      ,     )

						__SCREAMING_SNAKE_CASE				:    Any						    =							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 pytest

from datasets.splits import SplitDict, SplitInfo
from datasets.utils.py_utils import asdict


@pytest.mark.parametrize(
    '''split_dict''',   [
        SplitDict(),
        SplitDict({'''train''': SplitInfo(name='''train''',   num_bytes=1_337,   num_examples=42,   dataset_name='''my_dataset'''						)}						),
        SplitDict({'''train''': SplitInfo(name='''train''',   num_bytes=1_337,   num_examples=42						)}						),
        SplitDict({'''train''': SplitInfo()}						),
    ],   )
def        a__     (					UpperCamelCase_						:      SplitDict						):
   UpperCAmelCase__	:List[Any]      =      split_dict._to_yaml_list()
   assert len(UpperCamelCase_						) == len(UpperCamelCase_						)
   UpperCAmelCase__	:Tuple      =      SplitDict._from_yaml_list(UpperCamelCase_						)
   for split_name, split_info in split_dict.items():
      # dataset_name field is deprecated, and is therefore not part of the YAML dump
      UpperCAmelCase__	:Union[str, Any]      =      None
      # the split name of split_dict takes over the name of the split info object
      UpperCAmelCase__	:Union[str, Any]      =      split_name
   assert split_dict == reloaded


@pytest.mark.parametrize(
    '''split_info''',   [SplitInfo(), SplitInfo(dataset_name=UpperCamelCase_						), SplitInfo(dataset_name='''my_dataset'''						)]						)
def        a__     (					UpperCamelCase_						:      List[Any]						):
   # For backward compatibility, we need asdict(split_dict) to return split info dictrionaries with the "dataset_name"
   # field even if it's deprecated. This way old versionso of `datasets` can still reload dataset_infos.json files
   UpperCAmelCase__	:List[Any]      =      asdict(SplitDict({'''train''': split_info}						)						)
   assert "dataset_name" in split_dict_asdict["train"]
   assert split_dict_asdict["train"]["dataset_name"] == split_info.dataset_name



