





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        lowerCAmelCase__		(    self							):
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()




     @property
     def        lowerCAmelCase__		(    self							):
         _A     =							1
         _A     =							3
         _A     =							(32, 32)

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




     @property
     def        lowerCAmelCase__		(    self							):
         torch.manual_seed(0							)
         _A     =							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=snake_case_				,     only_cross_attention=(True, True, False)				,     num_class_embeds=100				,     )
         return model




     @property
     def        lowerCAmelCase__		(    self							):
         torch.manual_seed(0							)
         _A     =							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        lowerCAmelCase__		(    self							):
         torch.manual_seed(0							)
         _A     =							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(snake_case_							)




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

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

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

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

         _A     =							output.images

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

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

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

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

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

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

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




     @unittest.skipIf(torch_device != 'cuda'				,     'This test requires a GPU'							)
     def        lowerCAmelCase__		(    self							):
         _A     =							self.dummy_cond_unet_upscale
         _A     =							DDPMScheduler()
         _A     =							DDIMScheduler(prediction_type='v_prediction'							)
         _A     =							self.dummy_vae
         _A     =							self.dummy_text_encoder
         _A     =							CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'							)

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

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

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

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

         _A     =							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        lowerCAmelCase__		(    self							):
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()




     def        lowerCAmelCase__		(    self							):
         _A     =							load_image(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
             '/sd2-upscale/low_res_cat.png'							)
         _A     =							load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'
             '/upsampled_cat.npy'							)

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

         _A     =							'a cat sitting on a park bench'

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

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




     def        lowerCAmelCase__		(    self							):
         _A     =							load_image(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
             '/sd2-upscale/low_res_cat.png'							)
         _A     =							load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'
             '/upsampled_cat_fp16.npy'							)

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

         _A     =							'a cat sitting on a park bench'

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

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




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

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

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

         _A     =							'a cat sitting on a park bench'

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

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












import json
import os
from pathlib import Path

import pytest

from datasets.download.download_config import DownloadConfig
from datasets.download.download_manager import DownloadManager
from datasets.utils.file_utils import hash_url_to_filename


__A							:			List[Any]         =			"http://www.mocksite.com/file1.txt"
__A							:			List[Any]         =			"\"text\": [\"foo\", \"foo\"]"
__A							:			Dict         =			"6d8ce9aa78a471c7477201efbeabd3bb01ac2e7d100a6dc024ba1608361f90a8"







class      lowerCamelCase:
     '''simple docstring'''


     __magic_name__   =			200
     __magic_name__   =			{'Content-Length': '100'}
     __magic_name__   =			{}




     def        lowerCAmelCase__		(    self				,     **snake_case_							):
         return [bytes(snake_case_				,     'utf-8'							)]





def 				__lowerCAmelCase(       *_SCREAMING_SNAKE_CASE							,     **_SCREAMING_SNAKE_CASE				)							-> Optional[Any]:




    """simple docstring"""




    return MockResponse()





@pytest.mark.parametrize('urls_type'							,     [str, list, dict]				)
def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)							-> str:




    """simple docstring"""




    import requests

    monkeypatch.setattr(_SCREAMING_SNAKE_CASE							,     'request'							,     _SCREAMING_SNAKE_CASE				)

    _A     =							URL
    if issubclass(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
        _A     =							url
    elif issubclass(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
        _A     =							[url]
    elif issubclass(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
        _A     =							{'train': url}
    _A     =							'dummy'
    _A     =							'downloads'
    _A     =							tmp_path
    _A     =							DownloadConfig(
        cache_dir=os.path.join(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)							,     use_etag=_SCREAMING_SNAKE_CASE							,     )
    _A     =							DownloadManager(dataset_name=_SCREAMING_SNAKE_CASE							,     download_config=_SCREAMING_SNAKE_CASE				)
    _A     =							dl_manager.download(_SCREAMING_SNAKE_CASE				)
    _A     =							urls
    for downloaded_paths in [downloaded_paths]:
        if isinstance(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
            _A     =							[downloaded_paths]
            _A     =							[urls]
        elif isinstance(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
            assert "train" in downloaded_paths.keys()
            _A     =							downloaded_paths.values()
            _A     =							urls.values()
        assert downloaded_paths
        for downloaded_path, input_url in zip(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
            assert downloaded_path == dl_manager.downloaded_paths[input_url]
            _A     =							Path(_SCREAMING_SNAKE_CASE				)
            _A     =							downloaded_path.parts
            assert parts[-1] == HASH
            assert parts[-2] == cache_subdir
            assert downloaded_path.exists()
            _A     =							downloaded_path.read_text()
            assert content == CONTENT
            _A     =							downloaded_path.with_suffix('.json'				)
            assert metadata_downloaded_path.exists()
            _A     =							json.loads(metadata_downloaded_path.read_text()				)
            assert metadata_content == {"url": URL, "etag": None}





@pytest.mark.parametrize('paths_type'							,     [str, list, dict]				)
def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)							-> Optional[Any]:




    """simple docstring"""




    _A     =							str(_SCREAMING_SNAKE_CASE				)
    if issubclass(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
        _A     =							filename
    elif issubclass(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
        _A     =							[filename]
    elif issubclass(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
        _A     =							{'train': filename}
    _A     =							'dummy'
    _A     =							xz_file.parent
    _A     =							'extracted'
    _A     =							DownloadConfig(
        cache_dir=_SCREAMING_SNAKE_CASE							,     use_etag=_SCREAMING_SNAKE_CASE							,     )
    _A     =							DownloadManager(dataset_name=_SCREAMING_SNAKE_CASE							,     download_config=_SCREAMING_SNAKE_CASE				)
    _A     =							dl_manager.extract(_SCREAMING_SNAKE_CASE				)
    _A     =							paths
    for extracted_paths in [extracted_paths]:
        if isinstance(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
            _A     =							[extracted_paths]
            _A     =							[paths]
        elif isinstance(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
            assert "train" in extracted_paths.keys()
            _A     =							extracted_paths.values()
            _A     =							paths.values()
        assert extracted_paths
        for extracted_path, input_path in zip(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
            assert extracted_path == dl_manager.extracted_paths[input_path]
            _A     =							Path(_SCREAMING_SNAKE_CASE				)
            _A     =							extracted_path.parts
            assert parts[-1] == hash_url_to_filename(_SCREAMING_SNAKE_CASE							,     etag=_SCREAMING_SNAKE_CASE				)
            assert parts[-2] == extracted_subdir
            assert extracted_path.exists()
            _A     =							extracted_path.read_text()
            _A     =							text_file.read_text()
            assert extracted_file_content == expected_file_content





def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)							-> Dict:




    """simple docstring"""




    assert path.endswith('.jsonl'				)
    for num_items, line in enumerate(_SCREAMING_SNAKE_CASE							,     start=1				):
        _A     =							json.loads(line.decode('utf-8'				)				)
        assert item.keys() == {"col_1", "col_2", "col_3"}
    assert num_items == 4





@pytest.mark.parametrize('archive_jsonl'							,     ['tar_jsonl_path', 'zip_jsonl_path']				)
def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)							-> Optional[int]:




    """simple docstring"""




    _A     =							request.getfixturevalue(_SCREAMING_SNAKE_CASE				)
    _A     =							DownloadManager()
    for num_jsonl, (path, file) in enumerate(dl_manager.iter_archive(_SCREAMING_SNAKE_CASE				)							,     start=1				):
        _test_jsonl(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)
    assert num_jsonl == 2





@pytest.mark.parametrize('archive_nested_jsonl'							,     ['tar_nested_jsonl_path', 'zip_nested_jsonl_path']				)
def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)							-> Any:




    """simple docstring"""




    _A     =							request.getfixturevalue(_SCREAMING_SNAKE_CASE				)
    _A     =							DownloadManager()
    for num_tar, (path, file) in enumerate(dl_manager.iter_archive(_SCREAMING_SNAKE_CASE				)							,     start=1				):
        for num_jsonl, (subpath, subfile) in enumerate(dl_manager.iter_archive(_SCREAMING_SNAKE_CASE				)							,     start=1				):
            _test_jsonl(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)
    assert num_tar == 1
    assert num_jsonl == 2





def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE				)							-> int:




    """simple docstring"""




    _A     =							DownloadManager()
    for num_file, file in enumerate(dl_manager.iter_files(_SCREAMING_SNAKE_CASE				)							,     start=1				):
        assert os.path.basename(_SCREAMING_SNAKE_CASE				) == ("test.txt" if num_file == 1 else "train.txt")
    assert num_file == 2





