


from collections import Counter
from pathlib import Path
from typing import Optional, Tuple

import yaml






class      UpperCamelCase					(		yaml.SafeLoader     ):


   '''simple docstring'''







   def 		UpperCamelCase  (     self   ,							UpperCamelCase_		):
        lowercase_       :Dict     =       [self.constructed_objects[key_node] for key_node, _ in node.value]
        lowercase_       :List[Any]     =       [tuple(UpperCamelCase_		) if isinstance(UpperCamelCase_   ,							UpperCamelCase_		) else key for key in keys]
        lowercase_       :Tuple     =       Counter(UpperCamelCase_		)
        lowercase_       :Optional[int]     =       [key for key in counter if counter[key] > 1]
        if duplicate_keys:
             raise TypeError(f"Got duplicate yaml keys: {duplicate_keys}"		)




   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_=False		):
        lowercase_       :Dict     =       super().construct_mapping(UpperCamelCase_   ,							deep=UpperCamelCase_		)
        self._check_no_duplicates_on_constructed_node(UpperCamelCase_		)
        return mapping




def     UpperCamelCase    (      _a		)					->							Tuple[Optional[str], str]:
     '''simple docstring'''





     lowercase_       :Tuple     =       list(readme_content.splitlines()		)
     if full_content and full_content[0] == "---" and "---" in full_content[1:]:
          lowercase_       :List[str]     =       full_content[1:].index('''---'''		) + 1
          lowercase_       :Union[str, Any]     =       '''\n'''.join(full_content[1:sep_idx]		)
          return yamlblock, "\n".join(full_content[sep_idx + 1 :]		)

     return None, "\n".join(_a		)






class      UpperCamelCase					(		lowercase__     ):


   '''simple docstring'''







   lowercase					:	Any									={"""train_eval_index"""}  # train-eval-index in the YAML metadata
   @classmethod
   def 		UpperCamelCase  (     cls   ,							UpperCamelCase_		):
        with open(UpperCamelCase_   ,							encoding='''utf-8'''		) as readme_file:
             lowercase_    ,      lowercase_       :Union[str, Any]     =       _split_yaml_from_readme(readme_file.read()		)
        if yaml_string is not None:
             return cls.from_yaml_string(UpperCamelCase_		)
        else:
             return cls()
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_		):
        if path.exists():
             with open(UpperCamelCase_   ,							encoding='''utf-8'''		) as readme_file:
                  lowercase_       :Dict     =       readme_file.read()
        else:
             lowercase_       :List[str]     =       None
        lowercase_       :int     =       self._to_readme(UpperCamelCase_		)
        with open(UpperCamelCase_   ,							'''w'''   ,							encoding='''utf-8'''		) as readme_file:
             readme_file.write(UpperCamelCase_		)
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_ = None		):
        if readme_content is not None:
             lowercase_    ,      lowercase_       :List[Any]     =       _split_yaml_from_readme(UpperCamelCase_		)
             lowercase_       :Tuple     =       '''---\n''' + self.to_yaml_string() + '''---\n''' + content
        else:
             lowercase_       :Optional[int]     =       '''---\n''' + self.to_yaml_string() + '''---\n'''
        return full_content
   @classmethod
   def 		UpperCamelCase  (     cls   ,							UpperCamelCase_		):
        lowercase_       :str     =       yaml.load(UpperCamelCase_   ,							Loader=_NoDuplicateSafeLoader		) or {}

        # Convert the YAML keys to DatasetMetadata fields
        lowercase_       :Union[str, Any]     =       {
            (key.replace('''-'''   ,							'''_'''		) if key.replace('''-'''   ,							'''_'''		) in cls._FIELDS_WITH_DASHES else key): value
            for key, value in metadata_dict.items()
        }
        return cls(**UpperCamelCase_		)




   def 		UpperCamelCase  (     self		):
        return yaml.safe_dump(
            {
                (key.replace('''_'''   ,							'''-'''		) if key in self._FIELDS_WITH_DASHES else key): value
                for key, value in self.items()
            }   ,							sort_keys=UpperCamelCase_   ,							allow_unicode=UpperCamelCase_   ,							encoding='''utf-8'''   ,							).decode('''utf-8'''		)
SCREAMING_SNAKE_CASE       :		Any         =				{
    "image-classification": [],
    "translation": [],
    "image-segmentation": [],
    "fill-mask": [],
    "automatic-speech-recognition": [],
    "token-classification": [],
    "sentence-similarity": [],
    "audio-classification": [],
    "question-answering": [],
    "summarization": [],
    "zero-shot-classification": [],
    "table-to-text": [],
    "feature-extraction": [],
    "other": [],
    "multiple-choice": [],
    "text-classification": [],
    "text-to-image": [],
    "text2text-generation": [],
    "zero-shot-image-classification": [],
    "tabular-classification": [],
    "tabular-regression": [],
    "image-to-image": [],
    "tabular-to-text": [],
    "unconditional-image-generation": [],
    "text-retrieval": [],
    "text-to-speech": [],
    "object-detection": [],
    "audio-to-audio": [],
    "text-generation": [],
    "conversational": [],
    "table-question-answering": [],
    "visual-question-answering": [],
    "image-to-text": [],
    "reinforcement-learning": [],
    "voice-activity-detection": [],
    "time-series-forecasting": [],
    "document-question-answering": [],
}


if __name__ == "__main__":
   from argparse import ArgumentParser

   SCREAMING_SNAKE_CASE       :		Union[str, Any]         =				ArgumentParser(usage="Validate the yaml metadata block of a README.md file.")
   ap.add_argument("readme_filepath")
   SCREAMING_SNAKE_CASE       :		List[str]         =				ap.parse_args()

   SCREAMING_SNAKE_CASE       :		Dict         =				Path(args.readme_filepath)
   SCREAMING_SNAKE_CASE       :		int         =				DatasetMetadata.from_readme(readme_filepath)
   print(dataset_metadata)
   dataset_metadata.to_readme(readme_filepath)





import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
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      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   def 		UpperCamelCase  (     self		):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()
   @property
   def 		UpperCamelCase  (     self		):
        lowercase_       :Dict     =       1
        lowercase_       :Optional[Any]     =       3
        lowercase_       :Optional[int]     =       (32, 32)

        lowercase_       :Dict     =       floats_tensor((batch_size, num_channels) + sizes   ,							rng=random.Random(0		)		).to(UpperCamelCase_		)
        return image
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)
        lowercase_       :Optional[int]     =       UNetaDConditionModel(
            block_out_channels=(32, 64)   ,							layers_per_block=2   ,							sample_size=32   ,							in_channels=4   ,							out_channels=4   ,							down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')   ,							up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')   ,							cross_attention_dim=32   ,							)
        return model
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)
        lowercase_       :Dict     =       AutoencoderKL(
            block_out_channels=[32, 64]   ,							in_channels=3   ,							out_channels=3   ,							down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,							up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,							latent_channels=4   ,							)
        return model
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)
        lowercase_       :str     =       RobertaSeriesConfig(
            hidden_size=32   ,							project_dim=32   ,							intermediate_size=37   ,							layer_norm_eps=1E-05   ,							num_attention_heads=4   ,							num_hidden_layers=5   ,							pad_token_id=1   ,							vocab_size=5006   ,							)
        return RobertaSeriesModelWithTransformation(UpperCamelCase_		)
   @property
   def 		UpperCamelCase  (     self		):
        def extract(*UpperCamelCase_   ,							**UpperCamelCase_		):






             class      UpperCamelCase					:


                '''simple docstring'''







                def __init__(     self		):
                     lowercase_       :List[str]     =       torch.ones([0]		)




                def 		UpperCamelCase  (     self   ,							UpperCamelCase_		):
                     self.pixel_values.to(UpperCamelCase_		)
                     return self
             return Out()

        return extract
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[str]     =       '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        lowercase_       :Dict     =       self.dummy_cond_unet
        lowercase_       :Optional[Any]     =       PNDMScheduler(skip_prk_steps=UpperCamelCase_		)
        lowercase_       :Union[str, Any]     =       self.dummy_vae
        lowercase_       :Any     =       self.dummy_text_encoder
        lowercase_       :List[Any]     =       XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''		)
        lowercase_       :List[str]     =       77

        lowercase_       :int     =       self.dummy_image.to(UpperCamelCase_		)
        lowercase_       :Optional[Any]     =       init_image / 2 + 0.5

        # make sure here that pndm scheduler skips prk
        lowercase_       :str     =       AltDiffusionImgaImgPipeline(
            unet=UpperCamelCase_   ,							scheduler=UpperCamelCase_   ,							vae=UpperCamelCase_   ,							text_encoder=UpperCamelCase_   ,							tokenizer=UpperCamelCase_   ,							safety_checker=UpperCamelCase_   ,							feature_extractor=self.dummy_extractor   ,							)
        lowercase_       :List[str]     =       VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor   ,							do_normalize=UpperCamelCase_		)
        lowercase_       :int     =       alt_pipe.to(UpperCamelCase_		)
        alt_pipe.set_progress_bar_config(disable=UpperCamelCase_		)

        lowercase_       :Any     =       '''A painting of a squirrel eating a burger'''
        lowercase_       :Tuple     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :Optional[int]     =       alt_pipe(
            [prompt]   ,							generator=UpperCamelCase_   ,							guidance_scale=6.0   ,							num_inference_steps=2   ,							output_type='''np'''   ,							image=UpperCamelCase_   ,							)

        lowercase_       :Dict     =       output.images

        lowercase_       :Tuple     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :Optional[int]     =       alt_pipe(
            [prompt]   ,							generator=UpperCamelCase_   ,							guidance_scale=6.0   ,							num_inference_steps=2   ,							output_type='''np'''   ,							image=UpperCamelCase_   ,							return_dict=UpperCamelCase_   ,							)[0]

        lowercase_       :Optional[int]     =       image[0, -3:, -3:, -1]
        lowercase_       :Dict     =       image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 32, 32, 3)
        lowercase_       :int     =       np.array([0.4427, 0.3731, 0.4249, 0.4941, 0.4546, 0.4148, 0.4193, 0.4666, 0.4499]		)

        assert np.abs(image_slice.flatten() - expected_slice		).max() < 5E-3
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 5E-3
   @unittest.skipIf(torch_device != '''cuda'''   ,							'''This test requires a GPU'''		)
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[str]     =       self.dummy_cond_unet
        lowercase_       :Dict     =       PNDMScheduler(skip_prk_steps=UpperCamelCase_		)
        lowercase_       :Tuple     =       self.dummy_vae
        lowercase_       :Dict     =       self.dummy_text_encoder
        lowercase_       :Tuple     =       XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''		)
        lowercase_       :str     =       77

        lowercase_       :str     =       self.dummy_image.to(UpperCamelCase_		)

        # put models in fp16
        lowercase_       :Union[str, Any]     =       unet.half()
        lowercase_       :Union[str, Any]     =       vae.half()
        lowercase_       :List[str]     =       bert.half()

        # make sure here that pndm scheduler skips prk
        lowercase_       :List[Any]     =       AltDiffusionImgaImgPipeline(
            unet=UpperCamelCase_   ,							scheduler=UpperCamelCase_   ,							vae=UpperCamelCase_   ,							text_encoder=UpperCamelCase_   ,							tokenizer=UpperCamelCase_   ,							safety_checker=UpperCamelCase_   ,							feature_extractor=self.dummy_extractor   ,							)
        lowercase_       :List[Any]     =       VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor   ,							do_normalize=UpperCamelCase_		)
        lowercase_       :List[str]     =       alt_pipe.to(UpperCamelCase_		)
        alt_pipe.set_progress_bar_config(disable=UpperCamelCase_		)

        lowercase_       :str     =       '''A painting of a squirrel eating a burger'''
        lowercase_       :Union[str, Any]     =       torch.manual_seed(0		)
        lowercase_       :Any     =       alt_pipe(
            [prompt]   ,							generator=UpperCamelCase_   ,							num_inference_steps=2   ,							output_type='''np'''   ,							image=UpperCamelCase_   ,							).images

        assert image.shape == (1, 32, 32, 3)




   @unittest.skipIf(torch_device != '''cuda'''   ,							'''This test requires a GPU'''		)
   def 		UpperCamelCase  (     self		):
        lowercase_       :Optional[int]     =       load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/img2img/sketch-mountains-input.jpg'''		)
        # resize to resolution that is divisible by 8 but not 16 or 32
        lowercase_       :Optional[Any]     =       init_image.resize((760, 504)		)

        lowercase_       :List[str]     =       '''BAAI/AltDiffusion'''
        lowercase_       :Optional[Any]     =       AltDiffusionImgaImgPipeline.from_pretrained(
            UpperCamelCase_   ,							safety_checker=UpperCamelCase_   ,							)
        pipe.to(UpperCamelCase_		)
        pipe.set_progress_bar_config(disable=UpperCamelCase_		)
        pipe.enable_attention_slicing()

        lowercase_       :Optional[Any]     =       '''A fantasy landscape, trending on artstation'''

        lowercase_       :Optional[Any]     =       torch.manual_seed(0		)
        lowercase_       :str     =       pipe(
            prompt=UpperCamelCase_   ,							image=UpperCamelCase_   ,							strength=0.75   ,							guidance_scale=7.5   ,							generator=UpperCamelCase_   ,							output_type='''np'''   ,							)
        lowercase_       :Optional[Any]     =       output.images[0]

        lowercase_       :Optional[int]     =       image[255:258, 383:386, -1]

        assert image.shape == (504, 760, 3)
        lowercase_       :Any     =       np.array([0.9358, 0.9397, 0.9599, 0.9901, 1.0000, 1.0000, 0.9882, 1.0000, 1.0000]		)

        assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2






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


   '''simple docstring'''







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




   def 		UpperCamelCase  (     self		):
        lowercase_       :Union[str, Any]     =       load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/img2img/sketch-mountains-input.jpg'''		)
        lowercase_       :Any     =       init_image.resize((768, 512)		)
        lowercase_       :List[str]     =       load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy'''		)

        lowercase_       :List[Any]     =       '''BAAI/AltDiffusion'''
        lowercase_       :Tuple     =       AltDiffusionImgaImgPipeline.from_pretrained(
            UpperCamelCase_   ,							safety_checker=UpperCamelCase_   ,							)
        pipe.to(UpperCamelCase_		)
        pipe.set_progress_bar_config(disable=UpperCamelCase_		)
        pipe.enable_attention_slicing()

        lowercase_       :List[str]     =       '''A fantasy landscape, trending on artstation'''

        lowercase_       :Optional[int]     =       torch.manual_seed(0		)
        lowercase_       :Tuple     =       pipe(
            prompt=UpperCamelCase_   ,							image=UpperCamelCase_   ,							strength=0.75   ,							guidance_scale=7.5   ,							generator=UpperCamelCase_   ,							output_type='''np'''   ,							)
        lowercase_       :int     =       output.images[0]

        assert image.shape == (512, 768, 3)
        # img2img is flaky across GPUs even in fp32, so using MAE here
        assert np.abs(expected_image - image		).max() < 1E-2

