






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     (  _lowerCAmelCase ,      _lowerCAmelCase ,      _lowerCAmelCase ,      unittest.TestCase ):
	__snake_case						:Dict =							StableUnCLIPPipeline
	__snake_case						:List[Any] =							TEXT_TO_IMAGE_PARAMS
	__snake_case						:Dict =							TEXT_TO_IMAGE_BATCH_PARAMS
	__snake_case						:Union[str, Any] =							TEXT_TO_IMAGE_IMAGE_PARAMS
	__snake_case						:Any =							TEXT_TO_IMAGE_IMAGE_PARAMS

	# TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
	__snake_case						:Tuple =							False







	def 			_a (	self   :				str      )		->			Any:
		"""simple docstring"""







		__lowercase   							=	32
		__lowercase   							=	embedder_hidden_size

		# prior components

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

		torch.manual_seed(0      )
		__lowercase   							=	CLIPTextModelWithProjection(
		    CLIPTextConfig(
		        bos_token_id=0							,					eos_token_id=2							,					hidden_size=_lowerCAmelCase							,					projection_dim=_lowerCAmelCase							,					intermediate_size=37							,					layer_norm_eps=1e-05							,					num_attention_heads=4							,					num_hidden_layers=5							,					pad_token_id=1							,					vocab_size=1000							,					)      )

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

		torch.manual_seed(0      )
		__lowercase   							=	DDPMScheduler(
		    variance_type="""fixed_small_log"""							,					prediction_type="""sample"""							,					num_train_timesteps=1000							,					clip_sample=_lowerCAmelCase							,					clip_sample_range=5.0							,					beta_schedule="""squaredcos_cap_v2"""							,					)

		# regular denoising components

		torch.manual_seed(0      )
		__lowercase   							=	StableUnCLIPImageNormalizer(embedding_dim=_lowerCAmelCase      )
		__lowercase   							=	DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""      )

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

		torch.manual_seed(0      )
		__lowercase   							=	CLIPTextModel(
		    CLIPTextConfig(
		        bos_token_id=0							,					eos_token_id=2							,					hidden_size=_lowerCAmelCase							,					projection_dim=32							,					intermediate_size=37							,					layer_norm_eps=1e-05							,					num_attention_heads=4							,					num_hidden_layers=5							,					pad_token_id=1							,					vocab_size=1000							,					)      )

		torch.manual_seed(0      )
		__lowercase   							=	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=_lowerCAmelCase							,					layers_per_block=1							,					upcast_attention=_lowerCAmelCase							,					use_linear_projection=_lowerCAmelCase							,					)

		torch.manual_seed(0      )
		__lowercase   							=	DDIMScheduler(
		    beta_schedule="""scaled_linear"""							,					beta_start=0.00_085							,					beta_end=0.012							,					prediction_type="""v_prediction"""							,					set_alpha_to_one=_lowerCAmelCase							,					steps_offset=1							,					)

		torch.manual_seed(0      )
		__lowercase   							=	AutoencoderKL()

		__lowercase   							=	{
		    # 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 			_a (	self   :				int							,					_lowerCAmelCase   :				str							,					_lowerCAmelCase   :				Union[str, Any]=0      )		->			List[Any]:
		"""simple docstring"""







		if str(_lowerCAmelCase      ).startswith("""mps"""      ):
			__lowercase   							=	torch.manual_seed(_lowerCAmelCase      )
		else:
			__lowercase   							=	torch.Generator(device=_lowerCAmelCase      ).manual_seed(_lowerCAmelCase      )
		__lowercase   							=	{
		    """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 			_a (	self   :				Tuple      )		->			Tuple:
		"""simple docstring"""







		__lowercase   							=	torch_device == """cpu"""

		self._test_attention_slicing_forward_pass(test_max_difference=_lowerCAmelCase      )







	def 			_a (	self   :				Optional[Any]      )		->			str:
		"""simple docstring"""







		__lowercase   							=	torch_device in ["""cpu""", """mps"""]

		self._test_inference_batch_single_identical(test_max_difference=_lowerCAmelCase      )


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







	def 			_a (	self   :				Dict      )		->			Dict:
		"""simple docstring"""







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







	def 			_a (	self   :				Tuple      )		->			Any:
		"""simple docstring"""







		__lowercase   							=	load_numpy(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"""      )

		__lowercase   							=	StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""							,					torch_dtype=torch.floataa      )
		pipe.to(_lowerCAmelCase      )
		pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
		# 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()

		__lowercase   							=	torch.Generator(device="""cpu"""      ).manual_seed(0      )
		__lowercase   							=	pipe("""anime turle"""							,					generator=_lowerCAmelCase							,					output_type="""np"""      )

		__lowercase   							=	output.images[0]

		assert image.shape == (768, 768, 3)

		assert_mean_pixel_difference(_lowerCAmelCase							,					_lowerCAmelCase      )







	def 			_a (	self   :				List[str]      )		->			int:
		"""simple docstring"""







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

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

		__lowercase   							=	pipe(
		    """anime turtle"""							,					prior_num_inference_steps=2							,					num_inference_steps=2							,					output_type="""np"""							,					)

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


"""simple docstring"""



import datasets

from .evaluate import evaluate


_UpperCamelCase								=		"\\n@article{hendrycks2021cuad,\n      title={CUAD: An Expert-Annotated NLP Dataset for Legal Contract Review},\n      author={Dan Hendrycks and Collin Burns and Anya Chen and Spencer Ball},\n      journal={arXiv preprint arXiv:2103.06268},\n      year={2021}\n}\n"

_UpperCamelCase								=		"\nThis metric wrap the official scoring script for version 1 of the Contract\nUnderstanding Atticus Dataset (CUAD).\nContract Understanding Atticus Dataset (CUAD) v1 is a corpus of more than 13,000 labels in 510\ncommercial legal contracts that have been manually labeled to identify 41 categories of important\nclauses that lawyers look for when reviewing contracts in connection with corporate transactions.\n"

_UpperCamelCase								=		"\nComputes CUAD scores (EM, F1, AUPR, Precision@80%Recall, and Precision@90%Recall).\nArgs:\n    predictions: List of question-answers dictionaries with the following key-values:\n        - 'id': id of the question-answer pair as given in the references (see below)\n        - 'prediction_text': list of possible texts for the answer, as a list of strings\n        depending on a threshold on the confidence probability of each prediction.\n    references: List of question-answers dictionaries with the following key-values:\n        - 'id': id of the question-answer pair (see above),\n        - 'answers': a Dict in the CUAD dataset format\n            {\n                'text': list of possible texts for the answer, as a list of strings\n                'answer_start': list of start positions for the answer, as a list of ints\n            }\n            Note that answer_start values are not taken into account to compute the metric.\nReturns:\n    'exact_match': Exact match (the normalized answer exactly match the gold answer)\n    'f1': The F-score of predicted tokens versus the gold answer\n    'aupr': Area Under the Precision-Recall curve\n    'prec_at_80_recall': Precision at 80% recall\n    'prec_at_90_recall': Precision at 90% recall\nExamples:\n    >>> predictions = [{'prediction_text': ['The seller:', 'The buyer/End-User: Shenzhen LOHAS Supply Chain Management Co., Ltd.'], 'id': 'LohaCompanyltd_20191209_F-1_EX-10.16_11917878_EX-10.16_Supply Agreement__Parties'}]\n    >>> references = [{'answers': {'answer_start': [143, 49], 'text': ['The seller:', 'The buyer/End-User: Shenzhen LOHAS Supply Chain Management Co., Ltd.']}, 'id': 'LohaCompanyltd_20191209_F-1_EX-10.16_11917878_EX-10.16_Supply Agreement__Parties'}]\n    >>> cuad_metric = datasets.load_metric(\"cuad\")\n    >>> results = cuad_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'exact_match': 100.0, 'f1': 100.0, 'aupr': 0.0, 'prec_at_80_recall': 1.0, 'prec_at_90_recall': 1.0}\n"
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION							,   _KWARGS_DESCRIPTION     )
class 						__UpperCAmelCase    (datasets.Metric     ):







			'''simple docstring'''





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										return datasets.MetricInfo(
										    description=_DESCRIPTION   ,      citation=_CITATION   ,      inputs_description=_KWARGS_DESCRIPTION   ,      features=datasets.Features(
										        {
										            """predictions""": {
										                """id""": datasets.Value("""string"""							),
										                """prediction_text""": datasets.features.Sequence(datasets.Value("""string"""							)							),
										            },
										            """references""": {
										                """id""": datasets.Value("""string"""							),
										                """answers""": datasets.features.Sequence(
										                    {
										                        """text""": datasets.Value("""string"""							),
										                        """answer_start""": datasets.Value("""int32"""							),
										                    }							),
										            },
										        }							)   ,      codebase_urls=["""https://www.atticusprojectai.org/cuad"""]   ,      reference_urls=["""https://www.atticusprojectai.org/cuad"""]   ,      )





			def 		lowerCamelCase					(			self   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__       :     str								=			{prediction["""id"""]: prediction["""prediction_text"""] for prediction in predictions}
										A__       :     Optional[Any]								=			[
										    {
										        """paragraphs""": [
										            {
										                """qas""": [
										                    {
										                        """answers""": [{"""text""": answer_text} for answer_text in ref["""answers"""]["""text"""]],
										                        """id""": ref["""id"""],
										                    }
										                    for ref in references
										                ]
										            }
										        ]
										    }
										]
										A__       :     Optional[Any]								=			evaluate(dataset=snake_case_   ,      predictions=snake_case_							)
										return score





