




import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()


class 							__UpperCamelCase			( __UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					unittest.TestCase  ):





				'''simple docstring'''



				__a							: List[Any]										=StableUnCLIPPipeline
				__a							: List[str]										=TEXT_TO_IMAGE_PARAMS
				__a							: Any										=TEXT_TO_IMAGE_BATCH_PARAMS
				__a							: Any										=TEXT_TO_IMAGE_IMAGE_PARAMS
				__a							: str										=TEXT_TO_IMAGE_IMAGE_PARAMS

				# TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
				__a							: Optional[int]										=False



				def       __snake_case		( self						):
											lowerCAmelCase   =		32
											lowerCAmelCase   =		embedder_hidden_size

											# prior components

											torch.manual_seed(0						)
											lowerCAmelCase   =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						)

											torch.manual_seed(0						)
											lowerCAmelCase   =		CLIPTextModelWithProjection(
											    CLIPTextConfig(
											        bos_token_id=0       , eos_token_id=2       , hidden_size=UpperCAmelCase_       , projection_dim=UpperCAmelCase_       , intermediate_size=37       , layer_norm_eps=1E-0_5       , num_attention_heads=4       , num_hidden_layers=5       , pad_token_id=1       , vocab_size=10_00       , )						)

											torch.manual_seed(0						)
											lowerCAmelCase   =		PriorTransformer(
											    num_attention_heads=2       , attention_head_dim=12       , embedding_dim=UpperCAmelCase_       , num_layers=1       , )

											torch.manual_seed(0						)
											lowerCAmelCase   =		DDPMScheduler(
											    variance_type='''fixed_small_log'''       , prediction_type='''sample'''       , num_train_timesteps=10_00       , clip_sample=UpperCAmelCase_       , clip_sample_range=5.0       , beta_schedule='''squaredcos_cap_v2'''       , )

											# regular denoising components

											torch.manual_seed(0						)
											lowerCAmelCase   =		StableUnCLIPImageNormalizer(embedding_dim=UpperCAmelCase_						)
											lowerCAmelCase   =		DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''						)

											torch.manual_seed(0						)
											lowerCAmelCase   =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						)

											torch.manual_seed(0						)
											lowerCAmelCase   =		CLIPTextModel(
											    CLIPTextConfig(
											        bos_token_id=0       , eos_token_id=2       , hidden_size=UpperCAmelCase_       , projection_dim=32       , intermediate_size=37       , layer_norm_eps=1E-0_5       , num_attention_heads=4       , num_hidden_layers=5       , pad_token_id=1       , vocab_size=10_00       , )						)

											torch.manual_seed(0						)
											lowerCAmelCase   =		UNetaDConditionModel(
											    sample_size=32       , in_channels=4       , out_channels=4       , down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''')       , up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''')       , block_out_channels=(32, 64)       , attention_head_dim=(2, 4)       , class_embed_type='''projection'''       , projection_class_embeddings_input_dim=embedder_projection_dim * 2       , cross_attention_dim=UpperCAmelCase_       , layers_per_block=1       , upcast_attention=UpperCAmelCase_       , use_linear_projection=UpperCAmelCase_       , )

											torch.manual_seed(0						)
											lowerCAmelCase   =		DDIMScheduler(
											    beta_schedule='''scaled_linear'''       , beta_start=0.00085       , beta_end=0.012       , prediction_type='''v_prediction'''       , set_alpha_to_one=UpperCAmelCase_       , steps_offset=1       , )

											torch.manual_seed(0						)
											lowerCAmelCase   =		AutoencoderKL()

											lowerCAmelCase   =		{
											    # prior components
											    '''prior_tokenizer''': prior_tokenizer,
											    '''prior_text_encoder''': prior_text_encoder,
											    '''prior''': prior,
											    '''prior_scheduler''': prior_scheduler,
											    # image noising components
											    '''image_normalizer''': image_normalizer,
											    '''image_noising_scheduler''': image_noising_scheduler,
											    # regular denoising components
											    '''tokenizer''': tokenizer,
											    '''text_encoder''': text_encoder,
											    '''unet''': unet,
											    '''scheduler''': scheduler,
											    '''vae''': vae,
											}

											return components



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_=0						):
											if str(UpperCAmelCase_						).startswith('''mps'''						):
																		lowerCAmelCase   =		torch.manual_seed(UpperCAmelCase_						)
											else:
																		lowerCAmelCase   =		torch.Generator(device=UpperCAmelCase_						).manual_seed(UpperCAmelCase_						)
											lowerCAmelCase   =		{
											    '''prompt''': '''A painting of a squirrel eating a burger''',
											    '''generator''': generator,
											    '''num_inference_steps''': 2,
											    '''prior_num_inference_steps''': 2,
											    '''output_type''': '''numpy''',
											}
											return inputs



				def       __snake_case		( self						):
											lowerCAmelCase   =		torch_device == '''cpu'''

											self._test_attention_slicing_forward_pass(test_max_difference=UpperCAmelCase_						)






				def       __snake_case		( self						):
											lowerCAmelCase   =		torch_device in ['''cpu''', '''mps''']

											self._test_inference_batch_single_identical(test_max_difference=UpperCAmelCase_						)




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





				'''simple docstring'''



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



				def       __snake_case		( self						):
											lowerCAmelCase   =		load_numpy(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'''						)

											lowerCAmelCase   =		StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''       , torch_dtype=torch.floataa						)
											pipe.to(UpperCAmelCase_						)
											pipe.set_progress_bar_config(disable=UpperCAmelCase_						)
											# stable unclip will oom when integration tests are run on a V100,
											# so turn on memory savings
											pipe.enable_attention_slicing()
											pipe.enable_sequential_cpu_offload()

											lowerCAmelCase   =		torch.Generator(device='''cpu'''						).manual_seed(0						)
											lowerCAmelCase   =		pipe('''anime turle'''       , generator=UpperCAmelCase_       , output_type='''np'''						)

											lowerCAmelCase   =		output.images[0]

											assert image.shape == (7_68, 7_68, 3)

											assert_mean_pixel_difference(UpperCAmelCase_       , UpperCAmelCase_						)






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

											lowerCAmelCase   =		StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''       , torch_dtype=torch.floataa						)
											lowerCAmelCase   =		pipe.to(UpperCAmelCase_						)
											pipe.set_progress_bar_config(disable=UpperCAmelCase_						)
											pipe.enable_attention_slicing()
											pipe.enable_sequential_cpu_offload()

											lowerCAmelCase   =		pipe(
											    '''anime turtle'''       , prior_num_inference_steps=2       , num_inference_steps=2       , output_type='''np'''       , )

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









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

import yaml


class 							__UpperCamelCase			( yaml.SafeLoader  ):





				'''simple docstring'''



				def       __snake_case		( self       , UpperCAmelCase_						):
											lowerCAmelCase   =		[self.constructed_objects[key_node] for key_node, _ in node.value]
											lowerCAmelCase   =		[tuple(UpperCAmelCase_						) if isinstance(UpperCAmelCase_       , UpperCAmelCase_						) else key for key in keys]
											lowerCAmelCase   =		Counter(UpperCAmelCase_						)
											lowerCAmelCase   =		[key for key in counter if counter[key] > 1]
											if duplicate_keys:
																		raise TypeError(F"""Got duplicate yaml keys: {duplicate_keys}"""						)






				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_=False						):
											lowerCAmelCase   =		super().construct_mapping(UpperCAmelCase_       , deep=UpperCAmelCase_						)
											self._check_no_duplicates_on_constructed_node(UpperCAmelCase_						)
											return mapping




def       UpperCAmelCase       (     _snake_case							):
							lowerCAmelCase   =		list(readme_content.splitlines()							)
							if full_content and full_content[0] == "---" and "---" in full_content[1:]:
														lowerCAmelCase   =		full_content[1:].index('''---'''							) + 1
														lowerCAmelCase   =		'''\n'''.join(full_content[1:sep_idx]							)
														return yamlblock, "\n".join(full_content[sep_idx + 1 :]							)

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


class 							__UpperCamelCase			( __UpperCAmelCase  ):





				'''simple docstring'''



				__a							: Any										={"""train_eval_index"""}  # train-eval-index in the YAML metadata



				@classmethod
				def       __snake_case		( cls       , UpperCAmelCase_						):
											with open(UpperCAmelCase_       , encoding='''utf-8'''						) as readme_file:
																		lowerCAmelCase   ,     lowerCAmelCase   =		_split_yaml_from_readme(readme_file.read()						)
											if yaml_string is not None:
																		return cls.from_yaml_string(UpperCAmelCase_						)
											else:
																		return cls()



				def       __snake_case		( self       , UpperCAmelCase_						):
											if path.exists():
																		with open(UpperCAmelCase_       , encoding='''utf-8'''						) as readme_file:
																									lowerCAmelCase   =		readme_file.read()
											else:
																		lowerCAmelCase   =		None
											lowerCAmelCase   =		self._to_readme(UpperCAmelCase_						)
											with open(UpperCAmelCase_       , '''w'''       , encoding='''utf-8'''						) as readme_file:
																		readme_file.write(UpperCAmelCase_						)



				def       __snake_case		( self       , UpperCAmelCase_ = None						):
											if readme_content is not None:
																		lowerCAmelCase   ,     lowerCAmelCase   =		_split_yaml_from_readme(UpperCAmelCase_						)
																		lowerCAmelCase   =		'''---\n''' + self.to_yaml_string() + '''---\n''' + content
											else:
																		lowerCAmelCase   =		'''---\n''' + self.to_yaml_string() + '''---\n'''
											return full_content



				@classmethod
				def       __snake_case		( cls       , UpperCAmelCase_						):
											lowerCAmelCase   =		yaml.load(UpperCAmelCase_       , Loader=_NoDuplicateSafeLoader						) or {}

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






				def       __snake_case		( 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'''						)




UpperCAmelCase_						={
    """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

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

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



