




"""simple docstring"""





import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import (
					    MODEL_FOR_PRETRAINING_MAPPING,
					    MegatronBertForCausalLM,
					    MegatronBertForMaskedLM,
					    MegatronBertForMultipleChoice,
					    MegatronBertForNextSentencePrediction,
					    MegatronBertForPreTraining,
					    MegatronBertForQuestionAnswering,
					    MegatronBertForSequenceClassification,
					    MegatronBertForTokenClassification,
					    MegatronBertModel,
					)



class      a__				:


			def __init__(      self      :    Optional[Any]      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    Tuple=13      ,  UpperCamelCase_      :    int=7      ,  UpperCamelCase_      :    Tuple=True      ,  UpperCamelCase_      :    Any=True      ,  UpperCamelCase_      :    Dict=True      ,  UpperCamelCase_      :    int=True      ,  UpperCamelCase_      :    List[str]=99      ,  UpperCamelCase_      :    Any=64      ,  UpperCamelCase_      :    str=32      ,  UpperCamelCase_      :    Optional[Any]=5      ,  UpperCamelCase_      :    List[Any]=4      ,  UpperCamelCase_      :    List[Any]=37      ,  UpperCamelCase_      :    Dict="gelu"      ,  UpperCamelCase_      :    Union[str, Any]=0.1      ,  UpperCamelCase_      :    Optional[int]=0.1      ,  UpperCamelCase_      :    str=512      ,  UpperCamelCase_      :    Optional[Any]=16      ,  UpperCamelCase_      :    str=2      ,  UpperCamelCase_      :    Optional[int]=0.02      ,  UpperCamelCase_      :    Union[str, Any]=3      ,  UpperCamelCase_      :    Union[str, Any]=4      ,  UpperCamelCase_      :    Union[str, Any]=None      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					parent
						__UpperCAmelCase	: Any		=					batch_size
						__UpperCAmelCase	: Optional[Any]		=					seq_length
						__UpperCAmelCase	: List[Any]		=					is_training
						__UpperCAmelCase	: Optional[int]		=					use_input_mask
						__UpperCAmelCase	: Union[str, Any]		=					use_token_type_ids
						__UpperCAmelCase	: Optional[int]		=					use_labels
						__UpperCAmelCase	: Union[str, Any]		=					vocab_size
						__UpperCAmelCase	: Optional[int]		=					hidden_size
						__UpperCAmelCase	: str		=					embedding_size
						__UpperCAmelCase	: int		=					num_hidden_layers
						__UpperCAmelCase	: Dict		=					num_attention_heads
						__UpperCAmelCase	: int		=					intermediate_size
						__UpperCAmelCase	: Dict		=					hidden_act
						__UpperCAmelCase	: int		=					hidden_dropout_prob
						__UpperCAmelCase	: List[Any]		=					attention_probs_dropout_prob
						__UpperCAmelCase	: Optional[int]		=					max_position_embeddings
						__UpperCAmelCase	: List[Any]		=					type_vocab_size
						__UpperCAmelCase	: Any		=					type_sequence_label_size
						__UpperCAmelCase	: Dict		=					initializer_range
						__UpperCAmelCase	: Optional[Any]		=					num_labels
						__UpperCAmelCase	: List[str]		=					num_choices
						__UpperCAmelCase	: int		=					scope


			def       a_     (      self      :    List[str]):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					ids_tensor([self.batch_size, self.seq_length]      ,  self.vocab_size)

						__UpperCAmelCase	: Tuple		=					None
						if self.use_input_mask:
									__UpperCAmelCase	: List[Any]		=					random_attention_mask([self.batch_size, self.seq_length])

						__UpperCAmelCase	: Optional[Any]		=					None
						if self.use_token_type_ids:
									__UpperCAmelCase	: List[Any]		=					ids_tensor([self.batch_size, self.seq_length]      ,  self.type_vocab_size)

						__UpperCAmelCase	: Optional[int]		=					None
						__UpperCAmelCase	: List[str]		=					None
						__UpperCAmelCase	: str		=					None
						if self.use_labels:
									__UpperCAmelCase	: Any		=					ids_tensor([self.batch_size]      ,  self.type_sequence_label_size)
									__UpperCAmelCase	: int		=					ids_tensor([self.batch_size, self.seq_length]      ,  self.num_labels)
									__UpperCAmelCase	: Any		=					ids_tensor([self.batch_size]      ,  self.num_choices)

						__UpperCAmelCase	: Dict		=					self.get_config()

						return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




						return MegatronBertConfig(
						    vocab_size=self.vocab_size      ,  hidden_size=self.hidden_size      ,  num_hidden_layers=self.num_hidden_layers      ,  num_attention_heads=self.num_attention_heads      ,  intermediate_size=self.intermediate_size      ,  embedding_size=self.embedding_size      ,  hidden_act=self.hidden_act      ,  hidden_dropout_prob=self.hidden_dropout_prob      ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,  max_position_embeddings=self.max_position_embeddings      ,  type_vocab_size=self.type_vocab_size      ,  is_decoder=_lowerCAmelCase      ,  initializer_range=self.initializer_range      ,  )


			def       a_     (      self      :    Optional[int]      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					MegatronBertModel(config=_lowerCAmelCase)
						model.to(_lowerCAmelCase)
						model.eval()
						__UpperCAmelCase	: Any		=					model(_lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase)
						__UpperCAmelCase	: str		=					model(_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase)
						__UpperCAmelCase	: Tuple		=					model(_lowerCAmelCase)

						self.parent.assertEqual(result.last_hidden_state.shape      ,  (self.batch_size, self.seq_length, self.hidden_size))
						self.parent.assertEqual(result.pooler_output.shape      ,  (self.batch_size, self.hidden_size))


			def       a_     (      self      :    Tuple      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    Any):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					MegatronBertForMaskedLM(config=_lowerCAmelCase)
						model.to(_lowerCAmelCase)
						model.eval()
						__UpperCAmelCase	: List[Any]		=					model(_lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  labels=_lowerCAmelCase)
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.vocab_size))


			def       a_     (      self      :    Tuple      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    Dict):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					MegatronBertForCausalLM(config=_lowerCAmelCase)
						model.to(_lowerCAmelCase)
						model.eval()
						__UpperCAmelCase	: Tuple		=					model(_lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  labels=_lowerCAmelCase)
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.vocab_size))


			def       a_     (      self      :    Optional[int]      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: Tuple		=					MegatronBertForNextSentencePrediction(config=_lowerCAmelCase)
						model.to(_lowerCAmelCase)
						model.eval()
						__UpperCAmelCase	: Optional[Any]		=					model(
						    _lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  labels=_lowerCAmelCase      ,  )
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, 2))


			def       a_     (      self      :    Optional[Any]      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					MegatronBertForPreTraining(config=_lowerCAmelCase)
						model.to(_lowerCAmelCase)
						model.eval()
						__UpperCAmelCase	: Dict		=					model(
						    _lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  labels=_lowerCAmelCase      ,  next_sentence_label=_lowerCAmelCase      ,  )
						self.parent.assertEqual(result.prediction_logits.shape      ,  (self.batch_size, self.seq_length, self.vocab_size))
						self.parent.assertEqual(result.seq_relationship_logits.shape      ,  (self.batch_size, 2))


			def       a_     (      self      :    Optional[Any]      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    List[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: List[Any]		=					MegatronBertForQuestionAnswering(config=_lowerCAmelCase)
						model.to(_lowerCAmelCase)
						model.eval()
						__UpperCAmelCase	: Optional[int]		=					model(
						    _lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  start_positions=_lowerCAmelCase      ,  end_positions=_lowerCAmelCase      ,  )
						self.parent.assertEqual(result.start_logits.shape      ,  (self.batch_size, self.seq_length))
						self.parent.assertEqual(result.end_logits.shape      ,  (self.batch_size, self.seq_length))


			def       a_     (      self      :    Tuple      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    Union[str, Any]):


						"""simple docstring"""




						__UpperCAmelCase	: List[str]		=					self.num_labels
						__UpperCAmelCase	: str		=					MegatronBertForSequenceClassification(_lowerCAmelCase)
						model.to(_lowerCAmelCase)
						model.eval()
						__UpperCAmelCase	: Optional[Any]		=					model(_lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  labels=_lowerCAmelCase)
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.num_labels))


			def       a_     (      self      :    List[str]      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    Any):


						"""simple docstring"""




						__UpperCAmelCase	: List[str]		=					self.num_labels
						__UpperCAmelCase	: List[str]		=					MegatronBertForTokenClassification(config=_lowerCAmelCase)
						model.to(_lowerCAmelCase)
						model.eval()
						__UpperCAmelCase	: Union[str, Any]		=					model(_lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  labels=_lowerCAmelCase)
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.num_labels))


			def       a_     (      self      :    str      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: Dict		=					self.num_choices
						__UpperCAmelCase	: Any		=					MegatronBertForMultipleChoice(config=_lowerCAmelCase)
						model.to(_lowerCAmelCase)
						model.eval()
						__UpperCAmelCase	: int		=					input_ids.unsqueeze(1).expand(-1      ,  self.num_choices      ,  -1).contiguous()
						__UpperCAmelCase	: List[Any]		=					token_type_ids.unsqueeze(1).expand(-1      ,  self.num_choices      ,  -1).contiguous()
						__UpperCAmelCase	: int		=					input_mask.unsqueeze(1).expand(-1      ,  self.num_choices      ,  -1).contiguous()
						__UpperCAmelCase	: Union[str, Any]		=					model(
						    _lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  labels=_lowerCAmelCase      ,  )
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.num_choices))


			def       a_     (      self      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					self.prepare_config_and_inputs()
						(
						    (
						    __UpperCAmelCase			
						)	,     (
						    __UpperCAmelCase			
						)	,     (
						    __UpperCAmelCase			
						)	,     (
						    __UpperCAmelCase			
						)	,     (
						    __UpperCAmelCase			
						)	,     (
						    __UpperCAmelCase			
						)	,     (
						    __UpperCAmelCase			
						)	,     			
						)	: str		=					config_and_inputs
						__UpperCAmelCase	: Dict		=					{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
						return config, inputs_dict






@require_torch
class      a__				(	__magic_name__							,    __magic_name__							,    unittest.TestCase   ):
			lowercase_    =	(
			    (
			        MegatronBertModel,
			        MegatronBertForMaskedLM,
			        MegatronBertForCausalLM,
			        MegatronBertForMultipleChoice,
			        MegatronBertForNextSentencePrediction,
			        MegatronBertForPreTraining,
			        MegatronBertForQuestionAnswering,
			        MegatronBertForSequenceClassification,
			        MegatronBertForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			lowercase_    =	(
			    {
			        "feature-extraction": MegatronBertModel,
			        "fill-mask": MegatronBertForMaskedLM,
			        "question-answering": MegatronBertForQuestionAnswering,
			        "text-classification": MegatronBertForSequenceClassification,
			        "text-generation": MegatronBertForCausalLM,
			        "token-classification": MegatronBertForTokenClassification,
			        "zero-shot": MegatronBertForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			lowercase_    =	True
			# test_resize_embeddings = False
			lowercase_    =	False


			def       a_     (      self      :    Dict      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    List[str]=False):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					super()._prepare_for_class(_lowerCAmelCase      ,  _lowerCAmelCase      ,  return_labels=_lowerCAmelCase)

						if return_labels:
									if model_class in get_values(_lowerCAmelCase):
												__UpperCAmelCase	: str		=					torch.zeros(
												    (self.model_tester.batch_size, self.model_tester.seq_length)      ,  dtype=torch.long      ,  device=_lowerCAmelCase)
												__UpperCAmelCase	: Optional[Any]		=					torch.zeros(
												    self.model_tester.batch_size      ,  dtype=torch.long      ,  device=_lowerCAmelCase)
						return inputs_dict


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




						__UpperCAmelCase	: str		=					MegatronBertModelTester(self)
						__UpperCAmelCase	: Optional[Any]		=					ConfigTester(self      ,  config_class=_lowerCAmelCase      ,  hidden_size=37)


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




						self.config_tester.run_common_tests()


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_model(*_lowerCAmelCase)


			def       a_     (      self      :    Optional[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[Any]		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_masked_lm(*_lowerCAmelCase)


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




						__UpperCAmelCase	: List[Any]		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*_lowerCAmelCase)


			def       a_     (      self      :    List[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*_lowerCAmelCase)


			def       a_     (      self      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: List[str]		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_pretraining(*_lowerCAmelCase)


			def       a_     (      self      :    Dict):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_question_answering(*_lowerCAmelCase)


			def       a_     (      self      :    Optional[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: str		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*_lowerCAmelCase)


			def       a_     (      self      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: Any		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_token_classification(*_lowerCAmelCase)






def 					_UpperCamelCase    (			UpperCamelCase							)			->			Optional[int]:


			"""simple docstring"""
			return torch.tensor(
			    __lowerCAmelCase				,       dtype=torch.long				,       device=__lowerCAmelCase				,       )


A   =       1E-4





@require_torch
@require_sentencepiece
@require_tokenizers
class      a__				(	unittest.TestCase   ):


			@slow
			@unittest.skip("Model is not available.")
			def       a_     (      self      :    List[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					"nvidia/megatron-bert-uncased-345m"
						if "MYDIR" in os.environ:
									__UpperCAmelCase	: Optional[Any]		=					os.path.join(os.environ["MYDIR"]      ,  _lowerCAmelCase)
						__UpperCAmelCase	: int		=					MegatronBertModel.from_pretrained(_lowerCAmelCase)
						model.to(_lowerCAmelCase)
						model.half()
						__UpperCAmelCase	: Any		=					_long_tensor([[101, 7110, 1005, 1056, 2023, 11333, 17413, 1029, 102]])
						with torch.no_grad():
									__UpperCAmelCase	: List[str]		=					model(_lowerCAmelCase)[0]
						__UpperCAmelCase	: List[Any]		=					torch.Size((1, 9, 1024))
						self.assertEqual(output.shape      ,  _lowerCAmelCase)

						__UpperCAmelCase	: List[Any]		=					[-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
						for ii in range(3):
									for jj in range(3):
												__UpperCAmelCase	: Tuple		=					output[0, ii, jj]
												__UpperCAmelCase	: List[Any]		=					expected[3 * ii + jj]
												__UpperCAmelCase	: List[Any]		=					"ii={} jj={} a={} b={}".format(_lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase)
												self.assertTrue(math.isclose(_lowerCAmelCase      ,  _lowerCAmelCase      ,  rel_tol=_lowerCAmelCase      ,  abs_tol=_lowerCAmelCase)      ,  msg=_lowerCAmelCase)








"""simple docstring"""





import gc
import unittest

import numpy as np
import torch
from torch.backends.cuda import sdp_kernel

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)
from diffusers.utils import randn_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_a, require_torch_gpu

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class      a__				(	__magic_name__							,    unittest.TestCase   ):
			lowercase_    =	ConsistencyModelPipeline
			lowercase_    =	UNCONDITIONAL_IMAGE_GENERATION_PARAMS
			lowercase_    =	UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

			# Override required_optional_params to remove num_images_per_prompt
			lowercase_    =	frozenset(
			    [
			        "num_inference_steps",
			        "generator",
			        "latents",
			        "output_type",
			        "return_dict",
			        "callback",
			        "callback_steps",
			    ]   )


			@property
			def       a_     (      self      :    Optional[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					UNetaDModel.from_pretrained(
						    "diffusers/consistency-models-test"      ,  subfolder="test_unet"      ,  )
						return unet


			@property
			def       a_     (      self      :    List[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Dict		=					UNetaDModel.from_pretrained(
						    "diffusers/consistency-models-test"      ,  subfolder="test_unet_class_cond"      ,  )
						return unet


			def       a_     (      self      :    Any      ,  UpperCamelCase_      :    int=False):


						"""simple docstring"""




						if class_cond:
									__UpperCAmelCase	: List[Any]		=					self.dummy_cond_unet
						else:
									__UpperCAmelCase	: Optional[int]		=					self.dummy_uncond_unet

						# Default to CM multistep sampler
						__UpperCAmelCase	: List[str]		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )

						__UpperCAmelCase	: Optional[int]		=					{
						    "unet": unet,
						    "scheduler": scheduler,
						}

						return components


			def       a_     (      self      :    List[str]      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    List[Any]=0):


						"""simple docstring"""




						if str(UpperCamelCase_).startswith("mps"):
									__UpperCAmelCase	: str		=					torch.manual_seed(UpperCamelCase_)
						else:
									__UpperCAmelCase	: Optional[Any]		=					torch.Generator(device=UpperCamelCase_).manual_seed(UpperCamelCase_)

						__UpperCAmelCase	: List[Any]		=					{
						    "batch_size": 1,
						    "num_inference_steps": None,
						    "timesteps": [22, 0],
						    "generator": generator,
						    "output_type": "np",
						}

						return inputs


			def       a_     (      self      :    int):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					"cpu"  # ensure determinism for the device-dependent torch.Generator
						__UpperCAmelCase	: Union[str, Any]		=					self.get_dummy_components()
						__UpperCAmelCase	: str		=					ConsistencyModelPipeline(**UpperCamelCase_)
						__UpperCAmelCase	: Any		=					pipe.to(UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Any		=					self.get_dummy_inputs(UpperCamelCase_)
						__UpperCAmelCase	: str		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 32, 32, 3)

						__UpperCAmelCase	: Dict		=					image[0, -3:, -3:, -1]
						__UpperCAmelCase	: List[Any]		=					np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004])

						assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




						__UpperCAmelCase	: List[Any]		=					"cpu"  # ensure determinism for the device-dependent torch.Generator
						__UpperCAmelCase	: Optional[int]		=					self.get_dummy_components(class_cond=UpperCamelCase_)
						__UpperCAmelCase	: Optional[Any]		=					ConsistencyModelPipeline(**UpperCamelCase_)
						__UpperCAmelCase	: str		=					pipe.to(UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Tuple		=					self.get_dummy_inputs(UpperCamelCase_)
						__UpperCAmelCase	: Dict		=					0
						__UpperCAmelCase	: int		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 32, 32, 3)

						__UpperCAmelCase	: Any		=					image[0, -3:, -3:, -1]
						__UpperCAmelCase	: List[Any]		=					np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004])

						assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3


			def       a_     (      self      :    Optional[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					"cpu"  # ensure determinism for the device-dependent torch.Generator
						__UpperCAmelCase	: Optional[int]		=					self.get_dummy_components()
						__UpperCAmelCase	: Optional[Any]		=					ConsistencyModelPipeline(**UpperCamelCase_)
						__UpperCAmelCase	: Optional[Any]		=					pipe.to(UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: List[str]		=					self.get_dummy_inputs(UpperCamelCase_)
						__UpperCAmelCase	: Dict		=					1
						__UpperCAmelCase	: int		=					None
						__UpperCAmelCase	: List[Any]		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 32, 32, 3)

						__UpperCAmelCase	: Tuple		=					image[0, -3:, -3:, -1]
						__UpperCAmelCase	: Union[str, Any]		=					np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987])

						assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3


			def       a_     (      self      :    List[str]):


						"""simple docstring"""




						__UpperCAmelCase	: Dict		=					"cpu"  # ensure determinism for the device-dependent torch.Generator
						__UpperCAmelCase	: List[str]		=					self.get_dummy_components(class_cond=UpperCamelCase_)
						__UpperCAmelCase	: Tuple		=					ConsistencyModelPipeline(**UpperCamelCase_)
						__UpperCAmelCase	: int		=					pipe.to(UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Tuple		=					self.get_dummy_inputs(UpperCamelCase_)
						__UpperCAmelCase	: Optional[int]		=					1
						__UpperCAmelCase	: Optional[int]		=					None
						__UpperCAmelCase	: Tuple		=					0
						__UpperCAmelCase	: Tuple		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 32, 32, 3)

						__UpperCAmelCase	: List[str]		=					image[0, -3:, -3:, -1]
						__UpperCAmelCase	: Dict		=					np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987])

						assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3






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


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




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


			def       a_     (      self      :    Union[str, Any]      ,  UpperCamelCase_      :    Optional[int]=0      ,  UpperCamelCase_      :    Tuple=False      ,  UpperCamelCase_      :    int="cpu"      ,  UpperCamelCase_      :    Any=torch.floataa      ,  UpperCamelCase_      :    List[str]=(1, 3, 64, 64)):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					torch.manual_seed(UpperCamelCase_)

						__UpperCAmelCase	: int		=					{
						    "num_inference_steps": None,
						    "timesteps": [22, 0],
						    "class_labels": 0,
						    "generator": generator,
						    "output_type": "np",
						}

						if get_fixed_latents:
									__UpperCAmelCase	: int		=					self.get_fixed_latents(seed=UpperCamelCase_      ,  device=UpperCamelCase_      ,  dtype=UpperCamelCase_      ,  shape=UpperCamelCase_)
									__UpperCAmelCase	: Optional[int]		=					latents

						return inputs


			def       a_     (      self      :    Union[str, Any]      ,  UpperCamelCase_      :    int=0      ,  UpperCamelCase_      :    Tuple="cpu"      ,  UpperCamelCase_      :    Tuple=torch.floataa      ,  UpperCamelCase_      :    Optional[Any]=(1, 3, 64, 64)):


						"""simple docstring"""




						if type(UpperCamelCase_) == str:
									__UpperCAmelCase	: Union[str, Any]		=					torch.device(UpperCamelCase_)
						__UpperCAmelCase	: Union[str, Any]		=					torch.Generator(device=UpperCamelCase_).manual_seed(UpperCamelCase_)
						__UpperCAmelCase	: Optional[int]		=					randn_tensor(UpperCamelCase_      ,  generator=UpperCamelCase_      ,  device=UpperCamelCase_      ,  dtype=UpperCamelCase_)
						return latents


			def       a_     (      self      :    List[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					UNetaDModel.from_pretrained("diffusers/consistency_models"      ,  subfolder="diffusers_cd_imagenet64_l2")
						__UpperCAmelCase	: Dict		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )
						__UpperCAmelCase	: Dict		=					ConsistencyModelPipeline(unet=UpperCamelCase_      ,  scheduler=UpperCamelCase_)
						pipe.to(torch_device=UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Dict		=					self.get_inputs()
						__UpperCAmelCase	: List[str]		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 64, 64, 3)

						__UpperCAmelCase	: List[Any]		=					image[0, -3:, -3:, -1]

						__UpperCAmelCase	: Union[str, Any]		=					np.array([0.0888, 0.0881, 0.0666, 0.0479, 0.0292, 0.0195, 0.0201, 0.0163, 0.0254])

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


			def       a_     (      self      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: List[str]		=					UNetaDModel.from_pretrained("diffusers/consistency_models"      ,  subfolder="diffusers_cd_imagenet64_l2")
						__UpperCAmelCase	: Optional[Any]		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )
						__UpperCAmelCase	: Any		=					ConsistencyModelPipeline(unet=UpperCamelCase_      ,  scheduler=UpperCamelCase_)
						pipe.to(torch_device=UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: int		=					self.get_inputs()
						__UpperCAmelCase	: str		=					1
						__UpperCAmelCase	: Union[str, Any]		=					None
						__UpperCAmelCase	: Tuple		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 64, 64, 3)

						__UpperCAmelCase	: str		=					image[0, -3:, -3:, -1]

						__UpperCAmelCase	: Union[str, Any]		=					np.array([0.0340, 0.0152, 0.0063, 0.0267, 0.0221, 0.0107, 0.0416, 0.0186, 0.0217])

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


			@require_torch_a
			def       a_     (      self      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					UNetaDModel.from_pretrained("diffusers/consistency_models"      ,  subfolder="diffusers_cd_imagenet64_l2")
						__UpperCAmelCase	: Optional[int]		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )
						__UpperCAmelCase	: int		=					ConsistencyModelPipeline(unet=UpperCamelCase_      ,  scheduler=UpperCamelCase_)
						pipe.to(torch_device=UpperCamelCase_      ,  torch_dtype=torch.floataa)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Tuple		=					self.get_inputs(get_fixed_latents=UpperCamelCase_      ,  device=UpperCamelCase_)
						# Ensure usage of flash attention in torch 2.0
						with sdp_kernel(enable_flash=UpperCamelCase_      ,  enable_math=UpperCamelCase_      ,  enable_mem_efficient=UpperCamelCase_):
									__UpperCAmelCase	: List[str]		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 64, 64, 3)

						__UpperCAmelCase	: List[str]		=					image[0, -3:, -3:, -1]

						__UpperCAmelCase	: Optional[int]		=					np.array([0.1875, 0.1428, 0.1289, 0.2151, 0.2092, 0.1477, 0.1877, 0.1641, 0.1353])

						assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3


			@require_torch_a
			def       a_     (      self      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: Tuple		=					UNetaDModel.from_pretrained("diffusers/consistency_models"      ,  subfolder="diffusers_cd_imagenet64_l2")
						__UpperCAmelCase	: Union[str, Any]		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )
						__UpperCAmelCase	: List[Any]		=					ConsistencyModelPipeline(unet=UpperCamelCase_      ,  scheduler=UpperCamelCase_)
						pipe.to(torch_device=UpperCamelCase_      ,  torch_dtype=torch.floataa)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Optional[int]		=					self.get_inputs(get_fixed_latents=UpperCamelCase_      ,  device=UpperCamelCase_)
						__UpperCAmelCase	: List[str]		=					1
						__UpperCAmelCase	: Any		=					None
						# Ensure usage of flash attention in torch 2.0
						with sdp_kernel(enable_flash=UpperCamelCase_      ,  enable_math=UpperCamelCase_      ,  enable_mem_efficient=UpperCamelCase_):
									__UpperCAmelCase	: List[str]		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 64, 64, 3)

						__UpperCAmelCase	: int		=					image[0, -3:, -3:, -1]

						__UpperCAmelCase	: List[str]		=					np.array([0.1663, 0.1948, 0.2275, 0.1680, 0.1204, 0.1245, 0.1858, 0.1338, 0.2095])

						assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3


