





import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()





class 				A_				(      UpperCAmelCase					,	unittest.TestCase	):







       """simple docstring"""





       SCREAMING_SNAKE_CASE_      : List[str]      =  DiTPipeline
       SCREAMING_SNAKE_CASE_      : Tuple      =  CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
       SCREAMING_SNAKE_CASE_      : List[Any]      =  PipelineTesterMixin.required_optional_params - {
           '''latents''',
           '''num_images_per_prompt''',
           '''callback''',
           '''callback_steps''',
       }
       SCREAMING_SNAKE_CASE_      : List[str]      =  CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
       SCREAMING_SNAKE_CASE_      : str      =  False







       def    __UpperCAmelCase (    self	:	Optional[int]  )						->					int:
         torch.manual_seed(0  )
         _lowercase							= TransformeraDModel(
             sample_size=16						,num_layers=2						,patch_size=4						,attention_head_dim=8						,num_attention_heads=2						,in_channels=4						,out_channels=8						,attention_bias=__A						,activation_fn='gelu-approximate'						,num_embeds_ada_norm=1000						,norm_type='ada_norm_zero'						,norm_elementwise_affine=__A						,)
         _lowercase							= AutoencoderKL()
         _lowercase							= DDIMScheduler()
         _lowercase							= {'transformer': transformer.eval(), 'vae': vae.eval(), 'scheduler': scheduler}
         return components







       def    __UpperCAmelCase (    self	:	Optional[int]						,__A	:	List[Any]						,__A	:	str=0  )						->					Union[str, Any]:
         if str(__A  ).startswith('mps'  ):
           _lowercase							= torch.manual_seed(__A  )
         else:
           _lowercase							= torch.Generator(device=__A  ).manual_seed(__A  )
         _lowercase							= {
             'class_labels': [1],
             'generator': generator,
             'num_inference_steps': 2,
             'output_type': 'numpy',
         }
         return inputs







       def    __UpperCAmelCase (    self	:	Optional[Any]  )						->					List[str]:
         _lowercase							= 'cpu'

         _lowercase							= self.get_dummy_components()
         _lowercase							= self.pipeline_class(**__A  )
         pipe.to(__A  )
         pipe.set_progress_bar_config(disable=__A  )

         _lowercase							= self.get_dummy_inputs(__A  )
         _lowercase							= pipe(**__A  ).images
         _lowercase							= image[0, -3:, -3:, -1]

         self.assertEqual(image.shape						,(1, 16, 16, 3)  )
         _lowercase							= np.array([0.2946, 0.6601, 0.4329, 0.3296, 0.4144, 0.5319, 0.7273, 0.5013, 0.4457]  )
         _lowercase							= np.abs(image_slice.flatten() - expected_slice  ).max()
         self.assertLessEqual(__A						,1e-3  )







       def    __UpperCAmelCase (    self	:	Any  )						->					Tuple:
         self._test_inference_batch_single_identical(relax_max_difference=__A						,expected_max_diff=1e-3  )







       @unittest.skipIf(
           torch_device != 'cuda' or not is_xformers_available()						,reason='XFormers attention is only available with CUDA and `xformers` installed'						,)
       def    __UpperCAmelCase (    self	:	List[Any]  )						->					Dict:
         self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3  )





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







       """simple docstring"""





       def    __UpperCAmelCase (    self	:	str  )						->					List[Any]:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()







       def    __UpperCAmelCase (    self	:	Optional[Any]  )						->					Dict:
         _lowercase							= torch.manual_seed(0  )

         _lowercase							= DiTPipeline.from_pretrained('facebook/DiT-XL-2-256'  )
         pipe.to('cuda'  )

         _lowercase							= ['vase', 'umbrella', 'white shark', 'white wolf']
         _lowercase							= pipe.get_label_ids(__A  )

         _lowercase							= pipe(__A						,generator=__A						,num_inference_steps=40						,output_type='np'  ).images

         for word, image in zip(__A						,__A  ):
           _lowercase							= load_numpy(
               F"""https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy"""  )
           assert np.abs((expected_image - image).max()  ) < 1e-2







       def    __UpperCAmelCase (    self	:	Union[str, Any]  )						->					List[str]:
         _lowercase							= DiTPipeline.from_pretrained('facebook/DiT-XL-2-512'  )
         _lowercase							= DPMSolverMultistepScheduler.from_config(pipe.scheduler.config  )
         pipe.to('cuda'  )

         _lowercase							= ['vase', 'umbrella']
         _lowercase							= pipe.get_label_ids(__A  )

         _lowercase							= torch.manual_seed(0  )
         _lowercase							= pipe(__A						,generator=__A						,num_inference_steps=25						,output_type='np'  ).images

         for word, image in zip(__A						,__A  ):
           _lowercase							= load_numpy(
               'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
               F"""/dit/{word}_512.npy"""  )

           assert np.abs((expected_image - image).max()  ) < 1e-1






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

import yaml





class 				A_				(      yaml.SafeLoader	):







       """simple docstring"""





       def    __UpperCAmelCase (    self	:	Dict						,__A	:	List[Any]  )						->					Any:
         _lowercase							= [self.constructed_objects[key_node] for key_node, _ in node.value]
         _lowercase							= [tuple(__A  ) if isinstance(__A						,__A  ) else key for key in keys]
         _lowercase							= Counter(__A  )
         _lowercase							= [key for key in counter if counter[key] > 1]
         if duplicate_keys:
           raise TypeError(F"""Got duplicate yaml keys: {duplicate_keys}"""  )







       def    __UpperCAmelCase (    self	:	Any						,__A	:	int						,__A	:	Dict=False  )						->					Union[str, Any]:
         _lowercase							= super().construct_mapping(__A						,deep=__A  )
         self._check_no_duplicates_on_constructed_node(__A  )
         return mapping

def 							SCREAMING_SNAKE_CASE__       (  snake_case__  :str     )      ->						Tuple[Optional[str], str]:
  _lowercase							= list(readme_content.splitlines()     )
  if full_content and full_content[0] == "---" and "---" in full_content[1:]:
    _lowercase							= full_content[1:].index('---'     ) + 1
    _lowercase							= '\n'.join(full_content[1:sep_idx]     )
    return yamlblock, "\n".join(full_content[sep_idx + 1 :]     )

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





class 				A_				(      UpperCAmelCase	):







       """simple docstring"""





       SCREAMING_SNAKE_CASE_      : Any      =  {'''train_eval_index'''}  # train-eval-index in the YAML metadata







       @classmethod
       def    __UpperCAmelCase (    cls	:	List[Any]						,__A	:	Path  )						->					"DatasetMetadata":
         with open(__A						,encoding='utf-8'  ) as readme_file:
           _lowercase						,		_lowercase							= _split_yaml_from_readme(readme_file.read()  )
         if yaml_string is not None:
           return cls.from_yaml_string(__A  )
         else:
           return cls()







       def    __UpperCAmelCase (    self	:	Union[str, Any]						,__A	:	Path  )						->					Optional[Any]:
         if path.exists():
           with open(__A						,encoding='utf-8'  ) as readme_file:
             _lowercase							= readme_file.read()
         else:
           _lowercase							= None
         _lowercase							= self._to_readme(__A  )
         with open(__A						,'w'						,encoding='utf-8'  ) as readme_file:
           readme_file.write(__A  )







       def    __UpperCAmelCase (    self	:	Tuple						,__A	:	Optional[str] = None  )						->					str:
         if readme_content is not None:
           _lowercase						,		_lowercase							= _split_yaml_from_readme(__A  )
           _lowercase							= '---\n' + self.to_yaml_string() + '---\n' + content
         else:
           _lowercase							= '---\n' + self.to_yaml_string() + '---\n'
         return full_content







       @classmethod
       def    __UpperCAmelCase (    cls	:	List[Any]						,__A	:	str  )						->					"DatasetMetadata":
         _lowercase							= yaml.load(__A						,Loader=_NoDuplicateSafeLoader  ) or {}

         # Convert the YAML keys to DatasetMetadata fields
         _lowercase							= {
             (key.replace('-'						,'_'  ) if key.replace('-'						,'_'  ) in cls._FIELDS_WITH_DASHES else key): value
             for key, value in metadata_dict.items()
         }
         return cls(**__A  )







       def    __UpperCAmelCase (    self	:	Optional[int]  )						->					str:
         return yaml.safe_dump(
             {
                 (key.replace('_'						,'-'  ) if key in self._FIELDS_WITH_DASHES else key): value
                 for key, value in self.items()
             }						,sort_keys=__A						,allow_unicode=__A						,encoding='utf-8'						,).decode('utf-8'  )
snake_case 				=       {
    """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

      snake_case 				=       ArgumentParser(usage="""Validate the yaml metadata block of a README.md file.""")
      ap.add_argument("""readme_filepath""")
      snake_case 				=       ap.parse_args()

      snake_case 				=       Path(args.readme_filepath)
      snake_case 				=       DatasetMetadata.from_readme(readme_filepath)
      print(dataset_metadata)
      dataset_metadata.to_readme(readme_filepath)