





"""simple docstring"""



import re
import tempfile
from pathlib import Path

import pytest
import yaml

from datasets.utils.readme import ReadMe


# @pytest.fixture
# def example_yaml_structure():

a_ =			yaml.safe_load(
    """\
name: \"\"
allow_empty: false
allow_empty_text: true
subsections:
  - name: \"Dataset Card for X\" # First-level markdown heading
    allow_empty: false
    allow_empty_text: true
    subsections:
      - name: \"Table of Contents\"
        allow_empty: false
        allow_empty_text: false
        subsections: null
      - name: \"Dataset Description\"
        allow_empty: false
        allow_empty_text: false
        subsections:
          - name: \"Dataset Summary\"
            allow_empty: false
            allow_empty_text: false
            subsections: null
          - name: \"Supported Tasks and Leaderboards\"
            allow_empty: true
            allow_empty_text: true
            subsections: null
          - name: Languages
            allow_empty: false
            allow_empty_text: true
            subsections: null
"""
)


a_ =			{
    """name""": """root""",
    """text""": """""",
    """is_empty_text""": True,
    """subsections""": [
        {
            """name""": """Dataset Card for My Dataset""",
            """text""": """""",
            """is_empty_text""": True,
            """subsections""": [
                {"""name""": """Table of Contents""", """text""": """Some text here.""", """is_empty_text""": False, """subsections""": []},
                {
                    """name""": """Dataset Description""",
                    """text""": """Some text here.""",
                    """is_empty_text""": False,
                    """subsections""": [
                        {
                            """name""": """Dataset Summary""",
                            """text""": """Some text here.""",
                            """is_empty_text""": False,
                            """subsections""": [],
                        },
                        {
                            """name""": """Supported Tasks and Leaderboards""",
                            """text""": """""",
                            """is_empty_text""": True,
                            """subsections""": [],
                        },
                        {"""name""": """Languages""", """text""": """Language Text""", """is_empty_text""": False, """subsections""": []},
                    ],
                },
            ],
        }
    ],
}


a_ =			"""\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
"""


a_ =			"""\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
#### Extra Ignored Subsection
### Supported Tasks and Leaderboards
### Languages
Language Text
"""

a_ =			{
    """name""": """root""",
    """text""": """""",
    """is_empty_text""": True,
    """subsections""": [
        {
            """name""": """Dataset Card for My Dataset""",
            """text""": """""",
            """is_empty_text""": True,
            """subsections""": [
                {"""name""": """Table of Contents""", """text""": """Some text here.""", """is_empty_text""": False, """subsections""": []},
                {
                    """name""": """Dataset Description""",
                    """text""": """Some text here.""",
                    """is_empty_text""": False,
                    """subsections""": [
                        {
                            """name""": """Dataset Summary""",
                            """text""": """Some text here.""",
                            """is_empty_text""": False,
                            """subsections""": [
                                {
                                    """name""": """Extra Ignored Subsection""",
                                    """text""": """""",
                                    """is_empty_text""": True,
                                    """subsections""": [],
                                }
                            ],
                        },
                        {
                            """name""": """Supported Tasks and Leaderboards""",
                            """text""": """""",
                            """is_empty_text""": True,
                            """subsections""": [],
                        },
                        {"""name""": """Languages""", """text""": """Language Text""", """is_empty_text""": False, """subsections""": []},
                    ],
                },
            ],
        }
    ],
}

a_ =			"""\
---
---
# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
"""

a_ =			(
    """The following issues were found for the README at `{path}`:\n-\tEmpty YAML markers are present in the README."""
)

a_ =			"""\
# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
"""

a_ =			(
    """The following issues were found for the README at `{path}`:\n-\tNo YAML markers are present in the README."""
)

a_ =			"""\
---
# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
"""

a_ =			"""The following issues were found for the README at `{path}`:\n-\tOnly the start of YAML tags present in the README."""

a_ =			"""\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
### Supported Tasks and Leaderboards
### Languages
Language Text
"""
a_ =			"""The following issues were found for the README at `{path}`:\n-\tExpected some content in section `Dataset Summary` but it is empty.\n-\tExpected some text in section `Dataset Summary` but it is empty (text in subsections are ignored)."""


a_ =			"""\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
"""
a_ =			"""The following issues were found for the README at `{path}`:\n-\tExpected some content in section `Dataset Card for My Dataset` but it is empty.\n-\tSection `Dataset Card for My Dataset` expected the following subsections: `Table of Contents`, `Dataset Description`. Found 'None'."""

a_ =			"""\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Languages
Language Text
"""

a_ =			"""The following issues were found for the README at `{path}`:\n-\tSection `Dataset Description` is missing subsection: `Supported Tasks and Leaderboards`."""


a_ =			"""\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
"""

a_ =			"""The following issues were found for the README at `{path}`:\n-\tExpected some content in section `Languages` but it is empty."""

a_ =			"""\
---
language:
- zh
- en
---

## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
"""
a_ =			"""The following issues were found for the README at `{path}`:\n-\tThe README has no first-level headings. One heading is expected. Skipping further validation for this README."""

a_ =			"""\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
# Dataset Card My Dataset
"""

a_ =			"""The following issues were found for the README at `{path}`:\n-\tThe README has several first-level headings: `Dataset Card for My Dataset`, `Dataset Card My Dataset`. Only one heading is expected. Skipping further validation for this README."""

a_ =			"""\
---
language:
- zh
- en
---

# Dataset Card My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
"""

a_ =			"""The following issues were found for the README at `{path}`:\n-\tNo first-level heading starting with `Dataset Card for` found in README. Skipping further validation for this README."""

a_ =			""""""

a_ =			"""The following issues were found for the README at `{path}`:\n-\tThe README has no first-level headings. One heading is expected. Skipping further validation for this README.\n-\tNo YAML markers are present in the README."""

a_ =			"""\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
"""

a_ =			"""The following issues were found while parsing the README at `{path}`:\n-\tMultiple sections with the same heading `Dataset Card for My Dataset` have been found. Please keep only one of these sections."""







@pytest.mark.parametrize(
    '''readme_md, expected_dict'''				,[
        (README_CORRECT, CORRECT_DICT),
        (README_CORRECT_FOUR_LEVEL, CORRECT_DICT_FOUR_LEVEL),
    ]				,)
def 		__lowercase     (   snake_case_  :  Optional[int]				,snake_case_  :  List[Any]						)   ->Union[str, Any]:



    '''simple docstring'''







    assert ReadMe.from_string(snake_case_				,snake_case_						).to_dict() == expected_dict







@pytest.mark.parametrize(
    '''readme_md, expected_error'''				,[
        (README_NO_YAML, EXPECTED_ERROR_README_NO_YAML),
        (README_EMPTY_YAML, EXPECTED_ERROR_README_EMPTY_YAML),
        (README_INCORRECT_YAML, EXPECTED_ERROR_README_INCORRECT_YAML),
        (README_EMPTY, EXPECTED_ERROR_README_EMPTY),
        (README_NONE_SUBSECTION, EXPECTED_ERROR_README_NONE_SUBSECTION),
        (README_MISSING_FIRST_LEVEL, EXPECTED_ERROR_README_MISSING_FIRST_LEVEL),
        (README_MISSING_SUBSECTION, EXPECTED_ERROR_README_MISSING_SUBSECTION),
        (README_MISSING_TEXT, EXPECTED_ERROR_README_MISSING_TEXT),
        (README_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_WRONG_FIRST_LEVEL),
        (README_MULTIPLE_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_MULTIPLE_WRONG_FIRST_LEVEL),
        (README_MISSING_CONTENT, EXPECTED_ERROR_README_MISSING_CONTENT),
    ]				,)
def 		__lowercase     (   snake_case_  :  Dict				,snake_case_  :  Any						)   ->Dict:



    '''simple docstring'''







    with pytest.raises(snake_case_				,match=re.escape(expected_error.format(path='''root'''						)						)						):
        __A    :						str						 =     ReadMe.from_string(snake_case_				,snake_case_						)
        readme.validate()







@pytest.mark.parametrize(
    '''readme_md, expected_error'''				,[
        (README_MULTIPLE_SAME_HEADING_1, EXPECTED_ERROR_README_MULTIPLE_SAME_HEADING_1),
    ]				,)
def 		__lowercase     (   snake_case_  :  Union[str, Any]				,snake_case_  :  int						)   ->List[str]:



    '''simple docstring'''







    with pytest.raises(snake_case_				,match=re.escape(expected_error.format(path='''root'''						)						)						):
        ReadMe.from_string(snake_case_				,snake_case_						)







@pytest.mark.parametrize(
    '''readme_md,'''				,[
        (README_MULTIPLE_SAME_HEADING_1),
    ]				,)
def 		__lowercase     (   snake_case_  :  Tuple						)   ->Any:



    '''simple docstring'''







    ReadMe.from_string(snake_case_				,snake_case_				,suppress_parsing_errors=snake_case_						)







@pytest.mark.parametrize(
    '''readme_md, expected_dict'''				,[
        (README_CORRECT, CORRECT_DICT),
        (README_CORRECT_FOUR_LEVEL, CORRECT_DICT_FOUR_LEVEL),
    ]				,)
def 		__lowercase     (   snake_case_  :  Union[str, Any]				,snake_case_  :  Any						)   ->Tuple:



    '''simple docstring'''







    with tempfile.TemporaryDirectory() as tmp_dir:
        __A    :						Tuple						 =     Path(snake_case_						) / '''README.md'''
        with open(snake_case_				,'''w+'''						) as readme_file:
            readme_file.write(snake_case_						)
        __A    :						Dict						 =     ReadMe.from_readme(snake_case_				,snake_case_						).to_dict()
        assert out["name"] == path
        assert out["text"] == ""
        assert out["is_empty_text"]
        assert out["subsections"] == expected_dict["subsections"]







@pytest.mark.parametrize(
    '''readme_md, expected_error'''				,[
        (README_NO_YAML, EXPECTED_ERROR_README_NO_YAML),
        (README_EMPTY_YAML, EXPECTED_ERROR_README_EMPTY_YAML),
        (README_INCORRECT_YAML, EXPECTED_ERROR_README_INCORRECT_YAML),
        (README_EMPTY, EXPECTED_ERROR_README_EMPTY),
        (README_NONE_SUBSECTION, EXPECTED_ERROR_README_NONE_SUBSECTION),
        (README_MISSING_FIRST_LEVEL, EXPECTED_ERROR_README_MISSING_FIRST_LEVEL),
        (README_MISSING_SUBSECTION, EXPECTED_ERROR_README_MISSING_SUBSECTION),
        (README_MISSING_TEXT, EXPECTED_ERROR_README_MISSING_TEXT),
        (README_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_WRONG_FIRST_LEVEL),
        (README_MULTIPLE_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_MULTIPLE_WRONG_FIRST_LEVEL),
        (README_MISSING_CONTENT, EXPECTED_ERROR_README_MISSING_CONTENT),
    ]				,)
def 		__lowercase     (   snake_case_  :  Optional[Any]				,snake_case_  :  Dict						)   ->Optional[int]:



    '''simple docstring'''







    with tempfile.TemporaryDirectory() as tmp_dir:
        __A    :						Dict						 =     Path(snake_case_						) / '''README.md'''
        with open(snake_case_				,'''w+'''						) as readme_file:
            readme_file.write(snake_case_						)
        __A    :						int						 =     expected_error.format(path=snake_case_						)
        with pytest.raises(snake_case_				,match=re.escape(snake_case_						)						):
            __A    :						Any						 =     ReadMe.from_readme(snake_case_				,snake_case_						)
            readme.validate()







@pytest.mark.parametrize(
    '''readme_md, expected_error'''				,[
        (README_MULTIPLE_SAME_HEADING_1, EXPECTED_ERROR_README_MULTIPLE_SAME_HEADING_1),
    ]				,)
def 		__lowercase     (   snake_case_  :  Any				,snake_case_  :  int						)   ->Dict:



    '''simple docstring'''







    with tempfile.TemporaryDirectory() as tmp_dir:
        __A    :						Optional[Any]						 =     Path(snake_case_						) / '''README.md'''
        with open(snake_case_				,'''w+'''						) as readme_file:
            readme_file.write(snake_case_						)
        __A    :						List[str]						 =     expected_error.format(path=snake_case_						)
        with pytest.raises(snake_case_				,match=re.escape(snake_case_						)						):
            ReadMe.from_readme(snake_case_				,snake_case_						)







@pytest.mark.parametrize(
    '''readme_md,'''				,[
        (README_MULTIPLE_SAME_HEADING_1),
    ]				,)
def 		__lowercase     (   snake_case_  :  int						)   ->str:



    '''simple docstring'''







    with tempfile.TemporaryDirectory() as tmp_dir:
        __A    :						str						 =     Path(snake_case_						) / '''README.md'''
        with open(snake_case_				,'''w+'''						) as readme_file:
            readme_file.write(snake_case_						)
        ReadMe.from_readme(snake_case_				,snake_case_				,suppress_parsing_errors=snake_case_						)











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





 """simple docstring"""


 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




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

 @property
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Union[str, Any]						 =     1
     __A    :						Any						 =     3
     __A    :						List[str]						 =     (32, 32)

     __A    :						List[str]						 =     floats_tensor((batch_size, num_channels) + sizes					,						rng=random.Random(0  )  ).to(__lowerCamelCase  )
     return image

 @property
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     torch.manual_seed(0  )
     __A    :						List[Any]						 =     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 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     torch.manual_seed(0  )
     __A    :						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 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     torch.manual_seed(0  )
     __A    :						Tuple						 =     CLIPTextConfig(
         bos_token_id=0					,						eos_token_id=2					,						hidden_size=32					,						intermediate_size=37					,						layer_norm_eps=1e-0_5					,						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 					UpperCamelCase__(    self  ):



     '''simple docstring'''




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

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

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

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

     __A    :						Union[str, Any]						 =     output.images

     __A    :						List[str]						 =     torch.Generator(device=__lowerCamelCase  ).manual_seed(0  )
     __A    :						str						 =     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]

     __A    :						Tuple						 =     image[0, -3:, -3:, -1]
     __A    :						int						 =     image_from_tuple[0, -3:, -3:, -1]

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



     '''simple docstring'''




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

     __A    :						List[Any]						 =     self.dummy_image.cpu().permute(0					,						2					,						3					,						1  )[0]
     __A    :						int						 =     Image.fromarray(np.uinta(__lowerCamelCase  )  ).convert('''RGB'''  ).resize((64, 64)  )

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

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

     __A    :						Optional[Any]						 =     torch.Generator(device=__lowerCamelCase  ).manual_seed(0  )
     __A    :						Any						 =     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'''					,						)
     __A    :						Union[str, Any]						 =     output.images
     assert image.shape[0] == 2







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



     '''simple docstring'''




     __A    :						List[Any]						 =     self.dummy_cond_unet_upscale
     __A    :						int						 =     DDPMScheduler()
     __A    :						List[Any]						 =     DDIMScheduler(prediction_type='''v_prediction'''  )
     __A    :						Optional[Any]						 =     self.dummy_vae
     __A    :						List[str]						 =     self.dummy_text_encoder
     __A    :						Dict						 =     CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )

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

     # put models in fp16, except vae as it overflows in fp16
     __A    :						Union[str, Any]						 =     unet.half()
     __A    :						Optional[int]						 =     text_encoder.half()

     # make sure here that pndm scheduler skips prk
     __A    :						Optional[int]						 =     StableDiffusionUpscalePipeline(
         unet=__lowerCamelCase					,						low_res_scheduler=__lowerCamelCase					,						scheduler=__lowerCamelCase					,						vae=__lowerCamelCase					,						text_encoder=__lowerCamelCase					,						tokenizer=__lowerCamelCase					,						max_noise_level=350					,						)
     __A    :						Union[str, Any]						 =     sd_pipe.to(__lowerCamelCase  )
     sd_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

     __A    :						Union[str, Any]						 =     '''A painting of a squirrel eating a burger'''
     __A    :						Optional[Any]						 =     torch.manual_seed(0  )
     __A    :						Tuple						 =     sd_pipe(
         [prompt]					,						image=__lowerCamelCase					,						generator=__lowerCamelCase					,						num_inference_steps=2					,						output_type='''np'''					,						).images

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






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





 """simple docstring"""


 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




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

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						List[Any]						 =     load_image(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
         '''/sd2-upscale/low_res_cat.png'''  )
     __A    :						Dict						 =     load_numpy(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
         '''/upsampled_cat.npy'''  )

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

     __A    :						Union[str, Any]						 =     '''a cat sitting on a park bench'''

     __A    :						Union[str, Any]						 =     torch.manual_seed(0  )
     __A    :						Optional[Any]						 =     pipe(
         prompt=__lowerCamelCase					,						image=__lowerCamelCase					,						generator=__lowerCamelCase					,						output_type='''np'''					,						)
     __A    :						List[str]						 =     output.images[0]

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

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Dict						 =     load_image(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
         '''/sd2-upscale/low_res_cat.png'''  )
     __A    :						List[Any]						 =     load_numpy(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
         '''/upsampled_cat_fp16.npy'''  )

     __A    :						Optional[int]						 =     '''stabilityai/stable-diffusion-x4-upscaler'''
     __A    :						Optional[int]						 =     StableDiffusionUpscalePipeline.from_pretrained(
         __lowerCamelCase					,						torch_dtype=torch.floataa					,						)
     pipe.to(__lowerCamelCase  )
     pipe.set_progress_bar_config(disable=__lowerCamelCase  )
     pipe.enable_attention_slicing()

     __A    :						Dict						 =     '''a cat sitting on a park bench'''

     __A    :						Any						 =     torch.manual_seed(0  )
     __A    :						Optional[int]						 =     pipe(
         prompt=__lowerCamelCase					,						image=__lowerCamelCase					,						generator=__lowerCamelCase					,						output_type='''np'''					,						)
     __A    :						Any						 =     output.images[0]

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







 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




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

     __A    :						Union[str, Any]						 =     load_image(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
         '''/sd2-upscale/low_res_cat.png'''  )

     __A    :						List[str]						 =     '''stabilityai/stable-diffusion-x4-upscaler'''
     __A    :						Dict						 =     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()

     __A    :						Tuple						 =     '''a cat sitting on a park bench'''

     __A    :						Tuple						 =     torch.manual_seed(0  )
     __A    :						List[str]						 =     pipe(
         prompt=__lowerCamelCase					,						image=__lowerCamelCase					,						generator=__lowerCamelCase					,						num_inference_steps=5					,						output_type='''np'''					,						)

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




